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

从视频到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`，修改以下路径：

```python
# 视频源目录
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

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

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

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

### 3. 上传到CVAT

```bash
# 上传所有机位
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-*`参数跳过已完成的步骤