hoist-dataset-pipeline/config.py

#!/usr/bin/env python3
"""配置文件 - 所有可调参数集中管理"""
from pathlib import Path

# ============================================================
# 路径配置
# ============================================================
# 视频源目录
VIDEO_DIR = Path('/home/dell/disks/6T/dataset/gangbaowei/per-cangzhuozhongtie20260702')
# 视频源目录名（用于从路径中解析机位和日期）
VIDEO_DIR_NAME = 'per-cangzhuozhongtie20260702'

# 中间数据目录（抽帧结果）
EXTRACTED_DIR = Path('/home/dell/disks/6T/dataset/gangbaowei/dataset')
EXTRACTED_IMG_DIR = EXTRACTED_DIR / 'images'
EXTRACTED_LBL_DIR = EXTRACTED_DIR / 'labels'

# 最终输出目录
OUTPUT_DIR = Path('/home/dell/disks/6T/dataset/gangbaowei/dataset_release')

# CVAT ZIP输出目录
CVAT_ZIP_DIR = Path('/home/dell/disks/6T/dataset/gangbaowei/cvat_zips_release')

# ============================================================
# 模型配置
# ============================================================
# YOLOv5源码目录
YOLOV5_DIR = '/home/dell/deeplearn/yolov5-7.0-hoist-ear'
# 模型权重路径
MODEL_PATH = '/home/dell/deeplearn/yolov5-7.0-hoist-ear/runs/train/hoist-ear-all3/weights/best.pt'
# 检测类别
CLASSES = {0: 'hoist', 1: 'ear', 2: 'ready', 3: 'person', 4: 'car'}
# CVAT导入类别（6类，包含connector占位）
CVAT_CLASSES = ['hoist', 'ear', 'ready', 'person', 'car', 'connector']

# ============================================================
# 参数配置
# ============================================================
# 图片尺寸
IMG_WIDTH, IMG_HEIGHT = 640, 480

# Step 2: Engagement筛选
MAX_PER_CAM = 1000          # 每个机位最终保留的图片数
SCORE_THRESH = 0.35         # Engagement评分阈值
SEQ_GAP = 180               # 时序去重间隔（帧号差）

# Step 3: Ready注入
READY_CONF_THRESH = 0.25    # Ready检测置信度阈值
READY_IOU_THRESH = 0.45     # Ready NMS IoU阈值
READY_EAR_IOU_THRESH = 0.5  # Ready与ear的IoU过滤阈值

# Step 4: 质量过滤
EAR_FLAT_ASPECT_THRESH = 2.0      # ear宽高比阈值（超过视为扁平误检）
READY_HOIST_RATIO_THRESH = 0.9    # ready/hoist面积比阈值（超过视为误检）

# Step 5: 高阈值重推理
REINFER_CONF = 0.5          # 重推理置信度阈值
REINFER_IOU = 0.45          # 重推理NMS IoU阈值

# Step 6: 质量排序
QUALITY_BRIGHTNESS_WEIGHT = 0.4   # 亮度权重
QUALITY_SHARPNESS_WEIGHT = 0.6    # 清晰度权重

# ============================================================
# CVAT配置
# ============================================================
CVAT_URL = 'http://192.168.1.237:8080'
CVAT_USERNAME = 'lsfoo'
CVAT_PASSWORD = 'lsfoo!23Qwe'
CVAT_PROJECT_NAME = '钢包位识别'
# CVAT label ID映射（从项目中获取）
CVAT_LABEL_MAP = {0: 7, 1: 8, 2: 9}  # hoist=7, ear=8, ready=9