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

- pipeline.py: 7步自动化流程（抽帧→筛选→标注→过滤→重推理→排序→导出）
- upload_cvat.py: CVAT API自动上传
- config.py: 集中配置管理
- README.md: 使用说明

README.md

@@ -0,0 +1,183 @@
+# 起重机挂钩数据集自动生成工具
+
+从视频到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-*`参数跳过已完成的步骤

config.py

@@ -0,0 +1,72 @@
+#!/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

pipeline.py

@@ -0,0 +1,626 @@
+#!/usr/bin/env python3
+"""
+起重机挂钩数据集生成全流程
+============================
+Step 1: 视频抽帧 → Step 2: Engagement筛选 → Step 3: Ready注入
+→ Step 4: 质量过滤 → Step 5: 高阈值重推理 → Step 6: 质量排序 → Step 7: CVAT导出
+
+用法:
+  python pipeline.py                          # 运行全部步骤
+  python pipeline.py --skip-extract           # 跳过抽帧
+  python pipeline.py --skip-extract --skip-select  # 跳过多个步骤
+"""
+import os
+import sys
+import re
+import time
+import shutil
+import zipfile
+import subprocess
+import numpy as np
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+
+from config import *
+
+
+# ============================================================
+# Step 1: 视频抽帧 (FFmpeg NVDEC硬件加速)
+# ============================================================
+def step1_extract():
+    print("=" * 60)
+    print("Step 1: 视频抽帧 (FFmpeg NVDEC)")
+    print("=" * 60)
+
+    if EXTRACTED_DIR.exists():
+        print(f"  已存在 {EXTRACTED_DIR}，跳过抽帧")
+        return
+
+    videos = []
+    for mp4 in VIDEO_DIR.rglob('*.mp4'):
+        parts = mp4.parts
+        try:
+            cam_idx = parts.index(VIDEO_DIR_NAME)
+            cam_id = parts[cam_idx + 1]
+            cam_name = f'cam{cam_id.zfill(2)}'
+            date = parts[cam_idx + 2]
+        except (IndexError, ValueError):
+            continue
+        videos.append((mp4, cam_name, date))
+
+    print(f"  找到 {len(videos)} 个视频")
+    EXTRACTED_DIR.mkdir(parents=True, exist_ok=True)
+    total_frames = 0
+
+    for i, (mp4, cam_name, date) in enumerate(videos, 1):
+        out_dir = EXTRACTED_DIR / cam_name
+        out_dir.mkdir(exist_ok=True)
+        stem = mp4.stem
+        pattern = str(out_dir / f'{cam_name}_{date}_{stem}_%06d.jpg')
+
+        cmd = ['ffmpeg', '-hide_banner', '-loglevel', 'error',
+               '-hwaccel', 'nvdec', '-c:v', 'hevc_cuvid',
+               '-i', str(mp4),
+               '-qmin', '1', '-qmax', '1', '-vsync', '0', pattern]
+
+        print(f"  [{i}/{len(videos)}] {cam_name}/{date}/{stem}...", end=' ', flush=True)
+        t0 = time.time()
+        subprocess.run(cmd, check=True)
+
+        n_frames = len(list(out_dir.glob(f'{cam_name}_{date}_{stem}_*.jpg')))
+        for j in range(1, n_frames + 1):
+            lbl = out_dir / f'{cam_name}_{date}_{stem}_{j:06d}.txt'
+            if not lbl.exists():
+                lbl.touch()
+
+        elapsed = time.time() - t0
+        total_frames += n_frames
+        print(f"{n_frames} frames in {elapsed:.0f}s")
+
+    print(f"\n  抽帧完成: {total_frames} 帧")
+
+
+# ============================================================
+# Step 2: Engagement评分 + 去重 + 排序选取
+# ============================================================
+def _score_and_pos(dets):
+    """评分：hoist+ear共现=0.7, 仅ear=0.4, 中心距离奖励, 多hoist惩罚"""
+    if len(dets) == 0:
+        return 0
+    has_hoist = np.any(dets[:, 0] == 0)
+    has_ear = np.any(dets[:, 0] == 1)
+    n_hoist = int(np.sum(dets[:, 0] == 0))
+
+    score = 0.7 if (has_hoist and has_ear) else (0.4 if has_ear else 0.0)
+
+    avg_xc = np.mean(dets[:, 1])
+    avg_yc = np.mean(dets[:, 2])
+    center_dist = ((avg_xc - 0.5)**2 + (avg_yc - 0.5)**2) ** 0.5
+    score += max(0, 0.3 - center_dist) * 0.5
+    if n_hoist > 1:
+        score -= (n_hoist - 1) * 0.1
+    return max(0, score)
+
+
+def _fix_duplicates(labels):
+    """去重：ear保留最小框，hoist保留最大框"""
+    for cls_id, keep_func in [(1, 'min'), (0, 'max')]:
+        cls_mask = labels[:, 0] == cls_id
+        if np.sum(cls_mask) <= 1:
+            continue
+        cls_dets = labels[cls_mask]
+        areas = cls_dets[:, 3] * cls_dets[:, 4]
+        keep_idx = np.argmin(areas) if keep_func == 'min' else np.argmax(areas)
+        keep_mask = np.ones(len(labels), dtype=bool)
+        cls_indices = np.where(cls_mask)[0]
+        for ri in range(len(cls_dets)):
+            if ri != keep_idx:
+                keep_mask[cls_indices[ri]] = False
+        labels = labels[keep_mask]
+    return labels
+
+
+def step2_select():
+    print("\n" + "=" * 60)
+    print("Step 2: Engagement评分 + 去重 + 选取engaged帧")
+    print("=" * 60)
+
+    if OUTPUT_DIR.exists():
+        shutil.rmtree(str(OUTPUT_DIR))
+
+    cameras = sorted([d.name for d in EXTRACTED_IMG_DIR.iterdir() if d.is_dir()])
+    all_kept = {}
+
+    for cam_name in cameras:
+        t0 = time.time()
+        lbl_dir = EXTRACTED_LBL_DIR / cam_name
+
+        # 读取所有标签
+        labels = {}
+        with os.scandir(str(lbl_dir)) as it:
+            for entry in it:
+                if not entry.name.endswith('.txt'):
+                    continue
+                with open(entry.path) as f:
+                    data = f.read().strip()
+                stem = entry.name[:-4]
+                if data:
+                    rows = []
+                    for line in data.split('\n'):
+                        parts = line.split()
+                        if len(parts) >= 5:
+                            rows.append([float(x) for x in parts[:5]])
+                    labels[stem] = np.array(rows) if rows else np.zeros((0, 5))
+                else:
+                    labels[stem] = np.zeros((0, 5))
+
+        # 评分
+        scores = {name: _score_and_pos(dets) for name, dets in labels.items()}
+
+        # 去重
+        for name, dets in labels.items():
+            if len(dets) > 0:
+                labels[name] = _fix_duplicates(dets)
+
+        # 按分数排序
+        ranked = sorted(scores.keys(), key=lambda n: scores[n], reverse=True)
+
+        # 筛选
+        selected = []
+        last_seq = -999
+        for name in ranked:
+            if scores[name] < SCORE_THRESH:
+                continue
+            dets = labels[name]
+            if not np.any((dets[:, 0] == 0) | (dets[:, 0] == 1)):
+                continue
+            m = re.search(r'_(\d{6})$', name)
+            seq = int(m.group(1)) if m else 0
+            if abs(seq - last_seq) < SEQ_GAP or seq == last_seq:
+                continue
+            selected.append(name)
+            last_seq = seq
+
+        # 创建输出目录 + 符号链接
+        out_img = OUTPUT_DIR / 'images' / cam_name
+        out_lbl = OUTPUT_DIR / 'labels' / cam_name
+        out_img.mkdir(parents=True, exist_ok=True)
+        out_lbl.mkdir(parents=True, exist_ok=True)
+
+        img_dir = EXTRACTED_IMG_DIR / cam_name
+        for name in selected:
+            si = img_dir / f'{name}.jpg'
+            sl = EXTRACTED_LBL_DIR / cam_name / f'{name}.txt'
+            di = out_img / f'czzt_{name}.jpg'
+            dl = out_lbl / f'czzt_{name}.txt'
+            if si.exists() and not di.exists():
+                os.symlink(str(si.resolve()), str(di))
+            if sl.exists() and not dl.exists():
+                os.symlink(str(sl.resolve()), str(dl))
+
+        all_kept[cam_name] = len(selected)
+        print(f"  {cam_name}: {len(selected)} images in {time.time() - t0:.1f}s")
+
+    total = sum(all_kept.values())
+    print(f"\n  Step 2 完成: {total} images")
+    return all_kept
+
+
+# ============================================================
+# Step 3: 注入 ready 标签 (YOLOv5推理)
+# ============================================================
+def _calc_iou(boxA, boxB):
+    x1 = max(boxA[0], boxB[0]); y1 = max(boxA[1], boxB[1])
+    x2 = min(boxA[2], boxB[2]); y2 = min(boxA[3], boxB[3])
+    inter = max(0, x2-x1) * max(0, y2-y1)
+    areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
+    areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
+    union = areaA + areaB - inter
+    return inter / union if union > 0 else 0
+
+
+def step3_inject_ready(cam_name):
+    import torch
+    import cv2
+
+    img_dir = OUTPUT_DIR / 'images' / cam_name
+    lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+    if not img_dir.exists():
+        return 0
+
+    images = sorted(img_dir.glob('czzt_*.jpg'))
+    if not images:
+        return 0
+
+    sys.path.insert(0, YOLOV5_DIR)
+    from models.common import DetectMultiBackend
+    from utils.general import non_max_suppression, scale_boxes
+    from utils.torch_utils import select_device
+
+    device = select_device('')
+    model = DetectMultiBackend(MODEL_PATH, device=device)
+    model.eval()
+
+    injected = 0
+    for img_path in images:
+        lbl_file = lbl_dir / f'{img_path.stem}.txt'
+
+        # 读取已有标签
+        existing = []
+        if lbl_file.exists():
+            with open(lbl_file) as f:
+                for line in f:
+                    parts = line.strip().split()
+                    if len(parts) >= 5:
+                        existing.append([int(parts[0])] + [float(x) for x in parts[1:]])
+
+        # 需要同时有hoist和ear才注入ready
+        if not (any(r[0] == 0 for r in existing) and any(r[0] == 1 for r in existing)):
+            continue
+
+        # 推理
+        img_bgr = cv2.imread(str(img_path))
+        if img_bgr is None:
+            continue
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        h0, w0 = img_rgb.shape[:2]
+        img_resized = cv2.resize(img_rgb, (IMG_WIDTH, IMG_HEIGHT))
+        img_tensor = img_resized.transpose(2, 0, 1)
+        img_tensor = np.ascontiguousarray(img_tensor)
+        img_tensor = torch.from_numpy(img_tensor).to(device).float() / 255.0
+        if img_tensor.ndimension() == 3:
+            img_tensor = img_tensor.unsqueeze(0)
+
+        pred = model(img_tensor)[0]
+        pred = non_max_suppression(pred, READY_CONF_THRESH, READY_IOU_THRESH, classes=[2])
+
+        if len(pred) == 0 or len(pred[0]) == 0:
+            continue
+
+        dets = pred[0]
+        dets[:, :4] = scale_boxes(img_tensor.shape[2:], dets[:, :4], (h0, w0)).round()
+        areas = (dets[:, 2] - dets[:, 0]) * (dets[:, 3] - dets[:, 1])
+
+        # 收集ear框用于IoU过滤
+        ear_boxes = []
+        for r in existing:
+            if r[0] == 1:
+                ex, ey, ew, eh = r[1], r[2], r[3], r[4]
+                ear_boxes.append([(ex-ew/2)*w0, (ey-eh/2)*h0, (ex+ew/2)*w0, (ey+eh/2)*h0])
+
+        # 找最大的不与ear重叠的ready
+        sorted_idx = torch.argsort(areas, descending=True)
+        best = None
+        for idx in sorted_idx:
+            det = dets[idx]
+            rbox = [det[0].item(), det[1].item(), det[2].item(), det[3].item()]
+            if all(_calc_iou(rbox, ebox) <= READY_EAR_IOU_THRESH for ebox in ear_boxes):
+                best = det
+                break
+
+        if best is None:
+            continue
+
+        xc = (best[0] + best[2]) / 2 / w0
+        yc = (best[1] + best[3]) / 2 / h0
+        w = (best[2] - best[0]) / w0
+        h = (best[3] - best[1]) / h0
+
+        with open(lbl_file, 'a') as f:
+            f.write(f'2 {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n')
+        injected += 1
+
+    return injected
+
+
+def step3_inject():
+    print("\n" + "=" * 60)
+    print("Step 3: 注入 ready 标签 (YOLOv5推理)")
+    print("=" * 60)
+
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+    total_injected = 0
+    for cam_name in cameras:
+        n = step3_inject_ready(cam_name)
+        total_injected += n
+        print(f"  {cam_name}: injected ready to {n} images")
+    print(f"\n  Step 3 完成: {total_injected} ready labels injected")
+
+
+# ============================================================
+# Step 4: 质量过滤
+# ============================================================
+def step4_filter():
+    print("\n" + "=" * 60)
+    print("Step 4: 质量过滤")
+    print("=" * 60)
+
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+    total_keep = 0
+
+    for cam_name in cameras:
+        img_dir = OUTPUT_DIR / 'images' / cam_name
+        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+        if not lbl_dir.exists():
+            continue
+
+        discard_ratio = discard_flat = discard_both = keep = 0
+
+        for lbl_file in list(lbl_dir.glob('czzt_*.txt')):
+            with open(lbl_file) as f:
+                lines = f.readlines()
+
+            hoist_areas, ready_areas, ear_aspects = [], [], []
+            for line in lines:
+                parts = line.strip().split()
+                if len(parts) < 5:
+                    continue
+                cls = int(parts[0])
+                w, h = float(parts[3]), float(parts[4])
+                if cls == 0:
+                    hoist_areas.append(w * h)
+                elif cls == 2:
+                    ready_areas.append(w * h)
+                elif cls == 1 and h > 0:
+                    ear_aspects.append(w / h)
+
+            fail_ratio = (hoist_areas and ready_areas and
+                          max(ready_areas) / max(hoist_areas) >= READY_HOIST_RATIO_THRESH)
+            fail_flat = (ear_aspects and all(a >= EAR_FLAT_ASPECT_THRESH for a in ear_aspects))
+
+            if fail_ratio and fail_flat:
+                discard_both += 1
+            elif fail_ratio:
+                discard_ratio += 1
+            elif fail_flat:
+                discard_flat += 1
+            else:
+                keep += 1
+                continue
+
+            img_file = img_dir / (lbl_file.stem + '.jpg')
+            if img_file.exists():
+                img_file.unlink()
+            lbl_file.unlink()
+
+        total_keep += keep
+        print(f"  {cam_name}: keep={keep}, discard_ratio={discard_ratio}, discard_flat={discard_flat}, discard_both={discard_both}")
+
+    print(f"\n  Step 4 完成: {total_keep} images after filtering")
+
+
+# ============================================================
+# Step 5: 重新YOLO推理 (高置信度阈值, 清除重叠框)
+# ============================================================
+def step5_reinfer():
+    import torch
+    import cv2
+
+    print("\n" + "=" * 60)
+    print(f"Step 5: 重新YOLO推理 (conf>={REINFER_CONF}, iou={REINFER_IOU})")
+    print("=" * 60)
+
+    sys.path.insert(0, YOLOV5_DIR)
+    from models.common import DetectMultiBackend
+    from utils.general import non_max_suppression, scale_boxes
+    from utils.torch_utils import select_device
+
+    device = select_device('')
+    model = DetectMultiBackend(MODEL_PATH, device=device)
+    model.eval()
+
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+    total_reinferred = 0
+
+    for cam_name in cameras:
+        img_dir = OUTPUT_DIR / 'images' / cam_name
+        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+        if not img_dir.exists():
+            continue
+
+        images = sorted(img_dir.glob('czzt_*.jpg'))
+        if not images:
+            continue
+
+        reinferred = 0
+        for img_path in images:
+            img_bgr = cv2.imread(str(img_path))
+            if img_bgr is None:
+                continue
+            img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+            h0, w0 = img_rgb.shape[:2]
+            img_resized = cv2.resize(img_rgb, (IMG_WIDTH, IMG_HEIGHT))
+            img_tensor = img_resized.transpose(2, 0, 1)
+            img_tensor = np.ascontiguousarray(img_tensor)
+            img_tensor = torch.from_numpy(img_tensor).to(device).float() / 255.0
+            if img_tensor.ndimension() == 3:
+                img_tensor = img_tensor.unsqueeze(0)
+
+            pred = model(img_tensor)[0]
+            pred = non_max_suppression(pred, REINFER_CONF, REINFER_IOU)
+
+            lbl_file = lbl_dir / f'{img_path.stem}.txt'
+
+            if len(pred) == 0 or len(pred[0]) == 0:
+                lbl_file.write_text('')
+                reinferred += 1
+                continue
+
+            dets = pred[0]
+            dets[:, :4] = scale_boxes(img_tensor.shape[2:], dets[:, :4], (h0, w0)).round()
+
+            with open(lbl_file, 'w') as f:
+                for det in dets:
+                    cls_id = int(det[5])
+                    if cls_id not in CLASSES:
+                        continue
+                    xc = (det[0] + det[2]) / 2 / w0
+                    yc = (det[1] + det[3]) / 2 / h0
+                    w = (det[2] - det[0]) / w0
+                    h = (det[3] - det[1]) / h0
+                    f.write(f'{cls_id} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n')
+            reinferred += 1
+
+        total_reinferred += reinferred
+        print(f"  {cam_name}: reinferred {reinferred} images")
+
+    print(f"\n  Step 5 完成: {total_reinferred} images reinferred")
+
+
+# ============================================================
+# Step 6: 质量排序取top MAX_PER_CAM
+# ============================================================
+def step6_select_top():
+    print("\n" + "=" * 60)
+    print(f"Step 6: 质量排序取top {MAX_PER_CAM}")
+    print("=" * 60)
+
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+
+    for cam_name in cameras:
+        img_dir = OUTPUT_DIR / 'images' / cam_name
+        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+        if not img_dir.exists():
+            continue
+
+        images = sorted(img_dir.glob('czzt_*.jpg'))
+        if len(images) <= MAX_PER_CAM:
+            print(f"  {cam_name}: {len(images)} images (<= {MAX_PER_CAM}, no trimming)")
+            continue
+
+        def compute_quality(img_path):
+            try:
+                import cv2
+                img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
+                if img is None:
+                    return 0
+                brightness = np.mean(img) / 255.0
+                laplacian = cv2.Laplacian(img, cv2.CV_64F)
+                sharpness = laplacian.var() / 1000.0
+                return brightness * QUALITY_BRIGHTNESS_WEIGHT + min(sharpness, 1.0) * QUALITY_SHARPNESS_WEIGHT
+            except Exception:
+                return 0
+
+        quality = {}
+        with ThreadPoolExecutor(max_workers=16) as pool:
+            for img_path in images:
+                quality[img_path.name] = compute_quality(img_path)
+
+        ranked = sorted(quality.keys(), key=lambda n: quality[n], reverse=True)
+        keep_names = set(ranked[:MAX_PER_CAM])
+
+        removed = 0
+        for img_path in images:
+            if img_path.name not in keep_names:
+                img_path.unlink()
+                lbl = lbl_dir / (img_path.stem + '.txt')
+                if lbl.exists():
+                    lbl.unlink()
+                removed += 1
+
+        print(f"  {cam_name}: kept {len(keep_names)}, removed {removed}")
+
+
+# ============================================================
+# Step 7: 生成 CVAT ZIP
+# ============================================================
+def step7_generate_cvat():
+    print("\n" + "=" * 60)
+    print("Step 7: 生成 CVAT ZIP")
+    print("=" * 60)
+
+    CVAT_ZIP_DIR.mkdir(parents=True, exist_ok=True)
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+
+    for cam_name in cameras:
+        img_dir = OUTPUT_DIR / 'images' / cam_name
+        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+
+        img_names = sorted([f.name for f in img_dir.glob('czzt_*.jpg')])
+        if not img_names:
+            print(f"  {cam_name}: 0 images, skipping")
+            continue
+
+        zip_path = CVAT_ZIP_DIR / f'{cam_name}.zip'
+        obj_names_str = '\n'.join(CVAT_CLASSES) + '\n'
+        obj_data = f'classes = {len(CVAT_CLASSES)}\ntrain = data/train.txt\nnames = data/obj.names\nbackup = backup/\n'
+        train_lines = '\n'.join(f'data/obj_train_data/data/obj_train_data/{n}' for n in img_names) + '\n'
+
+        with zipfile.ZipFile(str(zip_path), 'w', zipfile.ZIP_STORED) as zf:
+            zf.writestr('obj.names', obj_names_str)
+            zf.writestr('obj.data', obj_data)
+            zf.writestr('train.txt', train_lines)
+            for img_name in img_names:
+                zf.write(str(img_dir / img_name), f'obj_train_data/data/obj_train_data/{img_name}')
+                stem = img_name[:-4]
+                lbl_file = lbl_dir / f'{stem}.txt'
+                if lbl_file.exists():
+                    zf.write(str(lbl_file), f'obj_train_data/data/obj_train_data/{stem}.txt')
+
+        print(f"  {cam_name}: {len(img_names)} images -> {zip_path.name}")
+
+    print(f"\n  Step 7 完成: CVAT zips in {CVAT_ZIP_DIR}")
+
+
+# ============================================================
+# 最终统计
+# ============================================================
+def print_stats():
+    print("\n" + "=" * 60)
+    print("最终统计")
+    print("=" * 60)
+
+    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
+    total = 0
+    label_counts = {0: 0, 1: 0, 2: 0}
+
+    for cam_name in cameras:
+        img_dir = OUTPUT_DIR / 'images' / cam_name
+        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
+        n_img = len(list(img_dir.glob('czzt_*.jpg'))) if img_dir.exists() else 0
+        n_lbl = {0: 0, 1: 0, 2: 0}
+        if lbl_dir.exists():
+            for lf in lbl_dir.glob('czzt_*.txt'):
+                with open(lf) as f:
+                    for line in f:
+                        parts = line.strip().split()
+                        if len(parts) >= 5:
+                            c = int(parts[0])
+                            if c in n_lbl:
+                                n_lbl[c] += 1
+                                label_counts[c] += 1
+        total += n_img
+        print(f"  {cam_name}: {n_img} images, hoist={n_lbl[0]}, ear={n_lbl[1]}, ready={n_lbl[2]}")
+
+    print(f"\n  总计: {total} images")
+    print(f"  hoist(0): {label_counts[0]}")
+    print(f"  ear(1): {label_counts[1]}")
+    print(f"  ready(2): {label_counts[2]}")
+
+
+# ============================================================
+# Main
+# ============================================================
+def main():
+    skip = set(sys.argv[1:])
+
+    if '--skip-extract' not in skip:
+        step1_extract()
+    if '--skip-select' not in skip:
+        step2_select()
+    if '--skip-inject' not in skip:
+        step3_inject()
+    if '--skip-filter' not in skip:
+        step4_filter()
+    if '--skip-reinfer' not in skip:
+        step5_reinfer()
+    if '--skip-select-top' not in skip:
+        step6_select_top()
+    if '--skip-zip' not in skip:
+        step7_generate_cvat()
+
+    print_stats()
+
+
+if __name__ == '__main__':
+    main()

upload_cvat.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""
+CVAT自动上传脚本
+================
+将pipeline生成的ZIP数据集自动上传到CVAT标注平台。
+
+用法:
+  python upload_cvat.py                  # 上传所有机位
+  python upload_cvat.py --cameras cam01 cam02  # 只上传指定机位
+"""
+import os
+import sys
+import time
+import zipfile
+import argparse
+import requests
+from pathlib import Path
+
+from config import *
+
+
+class CVATClient:
+    def __init__(self, url, username, password):
+        self.url = url.rstrip('/')
+        self.session = requests.Session()
+        self.login(username, password)
+
+    def login(self, username, password):
+        self.session.get(f'{self.url}/api/v1/auth/login')
+        csrf = self.session.cookies.get('csrftoken', '')
+        resp = self.session.post(f'{self.url}/api/v1/auth/login', json={
+            'username': username, 'password': password
+        }, headers={'X-CSRFToken': csrf})
+        resp.raise_for_status()
+        self.csrf = self.session.cookies.get('csrftoken', '')
+        print(f"Login OK")
+
+    def _headers(self):
+        return {'X-CSRFToken': self.csrf}
+
+    def get_project_id(self, name):
+        resp = self.session.get(f'{self.url}/api/v1/projects', params={'search': name})
+        resp.raise_for_status()
+        for p in resp.json()['results']:
+            if p['name'] == name:
+                return p['id']
+        raise ValueError(f"Project '{name}' not found")
+
+    def get_label_map(self, project_id):
+        resp = self.session.get(f'{self.url}/api/v1/projects/{project_id}')
+        resp.raise_for_status()
+        labels = resp.json().get('labels', [])
+        return {l['name']: l['id'] for l in labels}
+
+    def list_tasks(self, project_id):
+        resp = self.session.get(f'{self.url}/api/v1/tasks', params={'project': project_id})
+        resp.raise_for_status()
+        return resp.json()['results']
+
+    def create_task(self, project_id, task_name):
+        resp = self.session.post(f'{self.url}/api/v1/tasks', json={
+            'name': task_name, 'project_id': project_id
+        }, headers=self._headers())
+        resp.raise_for_status()
+        task = resp.json()
+        print(f"  Created task: {task['id']} - {task_name}")
+        return task['id']
+
+    def delete_task(self, task_id):
+        resp = self.session.delete(f'{self.url}/api/v1/tasks/{task_id}', headers=self._headers())
+        return resp.status_code
+
+    def upload_data(self, task_id, zip_path):
+        with open(zip_path, 'rb') as f:
+            resp = self.session.post(
+                f'{self.url}/api/v1/tasks/{task_id}/data',
+                files={'client_files[0]': (zip_path.name, f, 'application/zip')},
+                data={'image_quality': '100'},
+                headers=self._headers()
+            )
+        resp.raise_for_status()
+        print(f"  Uploaded data")
+        return resp.json()
+
+    def upload_annotations(self, task_id, zip_path, class_names, label_map):
+        shapes = []
+        with zipfile.ZipFile(zip_path) as zf:
+            ann_files = sorted([n for n in zf.namelist() if n.endswith('.txt') and 'obj_train_data' in n])
+            for i, ann_file in enumerate(ann_files):
+                with zf.open(ann_file) as f:
+                    for line in f:
+                        parts = line.decode().strip().split()
+                        if len(parts) >= 5:
+                            cls_id = int(parts[0])
+                            if cls_id not in class_names:
+                                continue
+                            label_name = class_names[cls_id]
+                            if label_name not in label_map:
+                                continue
+                            xc, yc, w, h = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
+                            shapes.append({
+                                'frame': i,
+                                'label_id': label_map[label_name],
+                                'type': 'rectangle',
+                                'occluded': False,
+                                'points': [
+                                    round((xc - w/2) * IMG_WIDTH, 2),
+                                    round((yc - h/2) * IMG_HEIGHT, 2),
+                                    round((xc + w/2) * IMG_WIDTH, 2),
+                                    round((yc + h/2) * IMG_HEIGHT, 2)
+                                ],
+                                'group': 0,
+                                'attributes': []
+                            })
+
+        if not shapes:
+            print(f"  No annotations to upload")
+            return
+
+        body = {'version': 0, 'tags': [], 'shapes': shapes, 'tracks': []}
+        resp = self.session.put(
+            f'{self.url}/api/v1/tasks/{task_id}/annotations',
+            json=body, headers=self._headers()
+        )
+        if resp.status_code == 200:
+            print(f"  Uploaded {len(shapes)} annotations")
+        else:
+            print(f"  Warning: {resp.status_code} {resp.text[:200]}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Upload dataset to CVAT')
+    parser.add_argument('--cameras', nargs='+', help='Specific cameras to upload (e.g. cam01 cam02)')
+    parser.add_argument('--delete-test', action='store_true', help='Delete test tasks (czzt_*_reinfer)')
+    args = parser.parse_args()
+
+    print("Connecting to CVAT...")
+    client = CVATClient(CVAT_URL, CVAT_USERNAME, CVAT_PASSWORD)
+
+    project_id = client.get_project_id(CVAT_PROJECT_NAME)
+    label_map = client.get_label_map(project_id)
+    class_names = {0: 'hoist', 1: 'ear', 2: 'ready'}
+    print(f"Project: {CVAT_PROJECT_NAME} (id={project_id}), labels: {label_map}")
+
+    # 删除测试任务
+    if args.delete_test:
+        tasks = client.list_tasks(project_id)
+        for t in tasks:
+            if '_reinfer' in t['name']:
+                print(f"  Deleting test task: {t['id']} - {t['name']}")
+                client.delete_task(t['id'])
+                time.sleep(0.5)
+
+    # 确定要上传的机位
+    if args.cameras:
+        zips = [CVAT_ZIP_DIR / f'{cam}.zip' for cam in args.cameras]
+        zips = [z for z in zips if z.exists()]
+    else:
+        zips = sorted(CVAT_ZIP_DIR.glob('*.zip'))
+
+    print(f"\nFound {len(zips)} zips to upload\n")
+
+    for zip_path in zips:
+        cam_name = zip_path.stem
+        print(f"\n{'='*40}")
+        print(f"Processing {cam_name}...")
+        print(f"{'='*40}")
+
+        task_id = client.create_task(project_id, f'czzt_{cam_name}')
+        client.upload_data(task_id, zip_path)
+        time.sleep(1)
+        client.upload_annotations(task_id, zip_path, class_names, label_map)
+        time.sleep(0.5)
+
+    print(f"\n{'='*40}")
+    print("Done!")
+
+
+if __name__ == '__main__':
+    main()