98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
"""
|
|
Training script for YOLO enemy/boss detection model.
|
|
|
|
Usage:
|
|
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 run_training(args):
|
|
"""Run YOLO training. Called from main() or manage.py."""
|
|
from ultralytics import YOLO
|
|
|
|
model = YOLO(f"{args.model}.pt")
|
|
|
|
model.train(
|
|
data=args.data,
|
|
epochs=args.epochs,
|
|
imgsz=args.imgsz,
|
|
batch=args.batch,
|
|
device=args.device,
|
|
name=args.name,
|
|
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,
|
|
)
|
|
|
|
# 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)
|
|
|
|
# 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")
|
|
|
|
|
|
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__":
|
|
main()
|