升级半自动工作台

2026-09-11 22:56:03 +08:00
parent f1c664f622
commit 4d6cded0b0
13 changed files with 3269 additions and 231 deletions
@@ -0,0 +1,69 @@
+"""
+YOLOv8n 目标检测器训练脚本
+基于准备好的 YOLO 数据集微调极轻量级模型，用于高精召回 672x480 背景图中的全部图标
+"""
+
+import logging
+from pathlib import Path
+import torch
+from ultralytics import YOLO
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("train_detector")
+
+DATASET_YAML = Path("tcaptcha/data/yolo/dataset.yaml")
+MODEL_DIR = Path("tcaptcha/models")
+
+
+def train(epochs: int = 40, batch: int = 8, imgsz: int = 640):
+    if not DATASET_YAML.exists():
+        logger.error("数据集配置文件未找到: %s，请先运行 prepare_yolo.py", DATASET_YAML)
+        return
+
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+    # 选择加速设备
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+    logger.info("采用训练加速设备: %s", device)
+
+    # 加载预训练 YOLOv8n (极轻量，只有 3.2M 参数)
+    model = YOLO("yolov8n.pt")
+
+    logger.info("开始微调训练，轮数: %d, 批大小: %d, 分辨率: %d", epochs, batch, imgsz)
+
+    results = model.train(
+        data=str(DATASET_YAML.resolve()),
+        epochs=epochs,
+        batch=batch,
+        imgsz=imgsz,
+        device=device,
+        project=str(MODEL_DIR / "runs"),
+        name="detector_exp",
+        exist_ok=True,
+        workers=2,
+        augment=True,
+        degrees=45.0,  # 随机旋转增强抗噪能力
+        flipud=0.2,
+        fliplr=0.5,
+        mosaic=1.0,
+        verbose=True,
+    )
+
+    # 保存最佳权重到 models/yolo_detector.pt
+    best_pt = MODEL_DIR / "runs" / "detector_exp" / "weights" / "best.pt"
+    target_pt = MODEL_DIR / "yolo_detector.pt"
+    if best_pt.exists():
+        import shutil
+        shutil.copy(best_pt, target_pt)
+        logger.info("✓ 训练完成！最佳模型权重已保存至: %s", target_pt)
+    else:
+        logger.warning("未找到 best.pt 权重")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", type=int, default=35, help="训练轮数 (默认: 35)")
+    parser.add_argument("--batch", type=int, default=8, help="Batch size (默认: 8)")
+    args = parser.parse_args()
+    train(epochs=args.epochs, batch=args.batch)