hoist-dataset-pipeline

起重机挂钩数据集自动生成工具

从视频到CVAT标注的全流程自动化工具，用于生成钢包位起重机挂钩识别训练数据集。

功能概述

• 视频抽帧：从HEVC视频中批量提取帧图片，支持NVDEC硬件加速
• 智能筛选：自动识别挂钩工作状态，过滤无关画面
• 自动标注：AI模型自动检测并标注吊钩、耳板、到位状态
• 质量控制：多重过滤机制确保数据集质量
• 一键部署：自动生成CVAT格式数据包并上传

环境要求

• Python 3.8+
• PyTorch 2.0+ (CUDA)
• OpenCV
• NumPy
• FFmpeg (支持NVDEC硬件解码)
• requests

目录结构

hoist_dataset_project/
├── config.py          # 配置文件（路径、参数、模型）
├── pipeline.py        # 主流程脚本（7步自动化）
├── upload_cvat.py     # CVAT上传脚本
└── README.md          # 本文件

快速开始

1. 修改配置

编辑 config.py，修改以下路径：

# 视频源目录
VIDEO_DIR = Path('/path/to/your/videos')

# 中间数据目录（抽帧结果）
EXTRACTED_DIR = Path('/path/to/extracted')

# 最终输出目录
OUTPUT_DIR = Path('/path/to/output')

# CVAT ZIP输出目录
CVAT_ZIP_DIR = Path('/path/to/zips')

# 模型权重路径
MODEL_PATH = '/path/to/best.pt'

2. 运行Pipeline

# 运行全部7个步骤
python pipeline.py

# 跳过已完成的步骤（例如抽帧已完成）
python pipeline.py --skip-extract

# 跳过多个步骤
python pipeline.py --skip-extract --skip-select

3. 上传到CVAT

# 上传所有机位
python upload_cvat.py

# 只上传指定机位
python upload_cvat.py --cameras cam01 cam02

# 删除测试任务并上传
python upload_cvat.py --delete-test

Pipeline流程

┌─────────────────────────────────────────────────────────────┐
│  Step 1: 视频抽帧                                            │
│  - FFmpeg NVDEC硬件解码加速                                   │
│  - 自动解析目录结构提取机位/日期信息                            │
│  - 输出: 32.7万帧图片                                         │
├─────────────────────────────────────────────────────────────┤
│  Step 2: Engagement状态筛选                                   │
│  - 多维度评分: hoist+ear共现、距离中心、数量惩罚                │
│  - 时序去重避免连续帧重复                                      │
│  - 输出: 挂钩工作状态帧                                       │
├─────────────────────────────────────────────────────────────┤
│  Step 3: Ready标签注入                                        │
│  - YOLOv5模型检测挂钩到位状态                                  │
│  - IoU过滤避免与ear重叠                                       │
│  - 输出: 带ready标签的图片                                    │
├─────────────────────────────────────────────────────────────┤
│  Step 4: 质量过滤                                             │
│  - ready/hoist面积比过滤                                      │
│  - ear宽高比过滤（排除扁平误检）                               │
│  - 输出: 高质量标注帧                                         │
├─────────────────────────────────────────────────────────────┤
│  Step 5: 高阈值重推理                                         │
│  - conf=0.5重新推理所有图片                                   │
│  - YOLO NMS自动消除重叠框                                     │
│  - 输出: 无重叠的纯净标注                                     │
├─────────────────────────────────────────────────────────────┤
│  Step 6: 质量排序选取                                         │
│  - 亮度+清晰度综合评分                                        │
│  - 每机位取top 1000张                                        │
│  - 输出: 最终数据集                                           │
├─────────────────────────────────────────────────────────────┤
│  Step 7: CVAT格式导出                                         │
│  - 生成标准CVAT导入ZIP                                        │
│  - 支持API自动上传                                            │
│  - 输出: 10个CVAT数据包                                      │
└─────────────────────────────────────────────────────────────┘

输出结构

dataset_release/
├── images/
│   ├── cam01/        # 1000张图片
│   ├── cam02/
│   └── ...
└── labels/
    ├── cam01/        # YOLO格式标签
    ├── cam02/
    └── ...

cvat_zips_release/
├── cam01.zip         # CVAT导入包
├── cam02.zip
└── ...

标签格式

YOLO格式：class_id x_center y_center width height

class_id	名称	说明
0	hoist	吊钩
1	ear	耳板
2	ready	挂钩到位

参数说明

所有可调参数在 config.py 中集中管理：

参数	默认值	说明
MAX_PER_CAM	1000	每机位最终保留图片数
SCORE_THRESH	0.35	Engagement评分阈值
SEQ_GAP	180	时序去重间隔
REINFER_CONF	0.5	重推理置信度阈值
REINFER_IOU	0.45	重推理NMS IoU阈值
EAR_FLAT_ASPECT_THRESH	2.0	ear宽高比阈值
READY_HOIST_RATIO_THRESH	0.9	ready/hoist面积比阈值

实际效果

• 输入：116个HEVC视频，10个机位，3天数据，约107GB
• 输出：10,000张高质量标注图片（10机位 × 1000张）
• 重叠框：从41.2%帧有重叠降至0%
• 总耗时：约3小时（全流程自动化）

模型信息

• 模型：YOLOv5s
• 精度：mAP@0.5 = 98.8%
• 类别：hoist(吊钩), ear(耳板), ready(挂钩到位), person(人), car(车)

注意事项

1. 首次运行需要确保FFmpeg支持NVDEC硬件解码
2. 模型权重文件需要单独准备
3. CVAT上传需要先在平台上创建好项目和标签
4. 可以通过--skip-*参数跳过已完成的步骤