refactor

tools/python-detect/train.py

@@ -0,0 +1,58 @@
+"""
+Training script for YOLOv11n enemy detection model.
+
+Usage:
+    python train.py --data path/to/data.yaml --epochs 100
+
+Expects YOLO-format dataset with data.yaml pointing to train/val image directories.
+Export from Roboflow in "YOLOv11" format.
+"""
+
+import argparse
+import os
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train YOLOv11n enemy detector")
+    parser.add_argument("--data", required=True, help="Path to data.yaml")
+    parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
+    parser.add_argument("--imgsz", type=int, default=640, help="Image size")
+    parser.add_argument("--batch", type=int, default=16, help="Batch size")
+    parser.add_argument("--device", default="0", help="CUDA device (0, cpu)")
+    parser.add_argument("--name", default="enemy-v1", help="Run name")
+    args = parser.parse_args()
+
+    from ultralytics import YOLO
+
+    model = YOLO("yolo11n.pt")  # start from pretrained nano
+
+    model.train(
+        data=args.data,
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        device=args.device,
+        name=args.name,
+        patience=20,       # early stopping
+        save=True,
+        save_period=10,
+        plots=True,
+        verbose=True,
+    )
+
+    # Copy best weights to models directory
+    best_path = os.path.join("runs", "detect", args.name, "weights", "best.pt")
+    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_path = os.path.join(output_dir, f"{args.name}.pt")
+    if os.path.exists(best_path):
+        import shutil
+        shutil.copy2(best_path, output_path)
+        print(f"\nBest model copied to: {output_path}")
+    else:
+        print(f"\nWarning: {best_path} not found — check training output")
+
+
+if __name__ == "__main__":
+    main()