work on well of souls and yolo detection

tools/python-detect/train.py

@@ -1,30 +1,24 @@
 """
-Training script for YOLOv11n enemy detection model.
+Training script for YOLO enemy/boss detection model.
 
 Usage:
-    python train.py --data path/to/data.yaml --epochs 100
+    python train.py --data path/to/data.yaml --epochs 200
+    python train.py --data path/to/data.yaml --model yolo11m --imgsz 1280 --epochs 300
 
 Expects YOLO-format dataset with data.yaml pointing to train/val image directories.
 Export from Roboflow in "YOLOv11" format.
 """
 
 import argparse
+import glob
 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()
-
+def run_training(args):
+    """Run YOLO training. Called from main() or manage.py."""
     from ultralytics import YOLO
 
-    model = YOLO("yolo11n.pt")  # start from pretrained nano
+    model = YOLO(f"{args.model}.pt")
 
     model.train(
         data=args.data,
@@ -33,25 +27,71 @@ def main():
         batch=args.batch,
         device=args.device,
         name=args.name,
-        patience=20,       # early stopping
+        patience=30,
+
+        # Learning rate (fine-tuning pretrained, not from scratch)
+        lr0=0.001,
+        lrf=0.01,
+        cos_lr=True,
+        warmup_epochs=5,
+        weight_decay=0.001,
+
+        # Augmentation tuned for boss glow/morph effects
+        hsv_h=0.03,
+        hsv_s=0.8,
+        hsv_v=0.6,
+        scale=0.7,
+        translate=0.2,
+        degrees=5.0,
+        mixup=0.15,
+        close_mosaic=15,
+        erasing=0.3,
+
+        workers=0,          # avoid multiprocessing paging file issues on Windows
         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")
+    # Find best.pt — try the trainer's save_dir first, then scan runs/detect/
+    best_path = None
+    save_dir = getattr(model.trainer, "save_dir", None)
+    if save_dir:
+        candidate = os.path.join(str(save_dir), "weights", "best.pt")
+        if os.path.exists(candidate):
+            best_path = candidate
+    if not best_path:
+        run_base = os.path.join("runs", "detect")
+        candidates = sorted(glob.glob(os.path.join(run_base, f"{args.name}*", "weights", "best.pt")))
+        best_path = candidates[-1] if candidates else os.path.join(run_base, 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 boss is set (from manage.py), deploy as boss-{boss}.pt; otherwise use run name
+    boss = getattr(args, "boss", None)
+    model_filename = f"boss-{boss}.pt" if boss else f"{args.name}.pt"
+    output_path = os.path.join(output_dir, model_filename)
     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")
+        print(f"\nWarning: {best_path} not found -- check training output")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train YOLO enemy/boss detector")
+    parser.add_argument("--data", required=True, help="Path to data.yaml")
+    parser.add_argument("--model", default="yolo11s", help="YOLO model variant (yolo11n, yolo11s, yolo11m)")
+    parser.add_argument("--epochs", type=int, default=200, help="Training epochs")
+    parser.add_argument("--imgsz", type=int, default=1280, help="Image size")
+    parser.add_argument("--batch", type=int, default=8, 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()
+    run_training(args)
 
 
 if __name__ == "__main__":