升级半自动工作台

tcaptcha/prepare_yolo.py

@@ -0,0 +1,144 @@
+"""
+将标注的 55 组样本转换为标准 YOLO 目标检测数据集
+生成:
+  tcaptcha/data/yolo/images/train/
+  tcaptcha/data/yolo/images/val/
+  tcaptcha/data/yolo/labels/train/
+  tcaptcha/data/yolo/labels/val/
+  tcaptcha/data/yolo/dataset.yaml
+"""
+
+import json
+import logging
+import random
+import shutil
+from pathlib import Path
+import cv2
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("prepare_yolo")
+
+RAW_DIR = Path("tcaptcha/data/raw")
+LABEL_FILE = Path("tcaptcha/data/labels.jsonl")
+YOLO_DIR = Path("tcaptcha/data/yolo")
+
+IMG_W = 672
+IMG_H = 480
+BOX_SIZE = 54  # 图标大致大小 (约 54x54 像素)
+
+
+def find_all_icons_in_bg(bg_path: Path, labeled_points: list) -> list:
+    """
+    结合人工标注点与轮廓补充检测，找到大图中所有图标（包含 3 个目标 + 2~3 个干扰图标）
+    返回: [(cx, cy, w, h), ...] 绝对像素坐标
+    """
+    img = cv2.imread(str(bg_path))
+    if img is None:
+        return []
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+    # 提取暗色轮廓找干扰图标
+    edges = cv2.Canny(gray, 40, 120)
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
+    closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
+    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    all_boxes = []
+    # 1. 先加入人工标注的 3 个目标（绝对精准）
+    for p in labeled_points:
+        cx, cy = p["x"], p["y"]
+        all_boxes.append((cx, cy, BOX_SIZE, BOX_SIZE))
+
+    # 2. 补充检测未标注的干扰图标
+    for cnt in contours:
+        x, y, w, h = cv2.boundingRect(cnt)
+        if 26 <= w <= 85 and 26 <= h <= 85 and (w * h) >= 650:
+            cx, cy = x + w // 2, y + h // 2
+            # 必须远离已有图标
+            if all(np.hypot(cx - lx, cy - ly) > 42 for lx, ly, _, _ in all_boxes):
+                # 限制尺寸
+                bw = min(68, max(46, w + 8))
+                bh = min(68, max(46, h + 8))
+                all_boxes.append((cx, cy, bw, bh))
+
+    return all_boxes
+
+
+def prepare_dataset(val_ratio: float = 0.2, seed: int = 42):
+    random.seed(seed)
+    if not LABEL_FILE.exists():
+        logger.error("标签文件不存在: %s", LABEL_FILE)
+        return
+
+    # 清理并创建目录
+    if YOLO_DIR.exists():
+        shutil.rmtree(YOLO_DIR)
+
+    for split in ["train", "val"]:
+        (YOLO_DIR / "images" / split).mkdir(parents=True, exist_ok=True)
+        (YOLO_DIR / "labels" / split).mkdir(parents=True, exist_ok=True)
+
+    with open(LABEL_FILE, "r", encoding="utf-8") as f:
+        records = [json.loads(line) for line in f if line.strip()]
+
+    random.shuffle(records)
+    val_count = max(5, int(len(records) * val_ratio))
+    train_records = records[val_count:]
+    val_records = records[:val_count]
+
+    logger.info("总标注样本: %d (训练集: %d, 验证集: %d)", len(records), len(train_records), len(val_records))
+
+    def process_split(record_list, split_name):
+        total_boxes = 0
+        for r in record_list:
+            sid = r["sample_id"]
+            src_img = RAW_DIR / f"{sid}_bg.png"
+            if not src_img.exists():
+                continue
+
+            dst_img = YOLO_DIR / "images" / split_name / f"{sid}.png"
+            shutil.copy(src_img, dst_img)
+
+            # 提取全部图标坐标
+            boxes = find_all_icons_in_bg(src_img, r["points"])
+            total_boxes += len(boxes)
+
+            # 写入 YOLO 格式标注: class_id x_center y_center width height (归一化到 0~1)
+            dst_label = YOLO_DIR / "labels" / split_name / f"{sid}.txt"
+            lines = []
+            for cx, cy, bw, bh in boxes:
+                # 边界截断保护
+                x1 = max(0, cx - bw / 2)
+                y1 = max(0, cy - bh / 2)
+                x2 = min(IMG_W, cx + bw / 2)
+                y2 = min(IMG_H, cy + bh / 2)
+
+                norm_cx = (x1 + x2) / 2.0 / IMG_W
+                norm_cy = (y1 + y2) / 2.0 / IMG_H
+                norm_w = (x2 - x1) / IMG_W
+                norm_h = (y2 - y1) / IMG_H
+
+                lines.append(f"0 {norm_cx:.6f} {norm_cy:.6f} {norm_w:.6f} {norm_h:.6f}")
+
+            dst_label.write_text("\n".join(lines), encoding="utf-8")
+        logger.info("  [%s] 处理完成: %d 张图片, 共 %d 个图标标注", split_name, len(record_list), total_boxes)
+
+    process_split(train_records, "train")
+    process_split(val_records, "val")
+
+    # 生成 dataset.yaml
+    yaml_content = f"""path: {YOLO_DIR.resolve()}
+train: images/train
+val: images/val
+
+names:
+  0: icon
+"""
+    yaml_path = YOLO_DIR / "dataset.yaml"
+    yaml_path.write_text(yaml_content, encoding="utf-8")
+    logger.info("✓ YOLO 数据集生成完毕: %s", yaml_path)
+
+
+if __name__ == "__main__":
+    prepare_dataset()