58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
"""
|
|
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()
|