DirtyLeague_Bot/train_yolo.py

"""
YOLOv8 Training and ONNX Export Pipeline for DirtyLeague Bot.

1. Prepares train/val splits from dataset/raw and dataset/labels.
2. Trains YOLOv8-nano model on GPU (NVIDIA GTX 1060) or CPU.
3. Automatically exports the trained model to ONNX format (models/dirty_league_yolo.onnx).
"""

import os
import shutil
import random
import sys

DATASET_DIR = os.path.join(os.path.dirname(__file__), "dataset")
RAW_DIR = os.path.join(DATASET_DIR, "raw")
LABELS_DIR = os.path.join(DATASET_DIR, "labels")
MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")


def prepare_dataset_splits(val_ratio: float = 0.20):
    """Organizes images and labels into standard YOLO train/val folders.
    Ensures all classes are represented in train."""
    images_train = os.path.join(DATASET_DIR, "images", "train")
    images_val = os.path.join(DATASET_DIR, "images", "val")
    labels_train = os.path.join(DATASET_DIR, "labels", "train")
    labels_val = os.path.join(DATASET_DIR, "labels", "val")

    # Clean existing directories
    for d in [images_train, images_val, labels_train, labels_val]:
        if os.path.exists(d):
            shutil.rmtree(d)
        os.makedirs(d, exist_ok=True)

    image_files = [f for f in os.listdir(RAW_DIR) if f.endswith(".png")]
    random.seed(42)
    random.shuffle(image_files)

    # Detect rare classes to guarantee they exist in train
    rare_files = set()
    for fname in image_files:
        base_name = os.path.splitext(fname)[0]
        lbl_file = os.path.join(LABELS_DIR, f"{base_name}.txt")
        if os.path.exists(lbl_file):
            with open(lbl_file, "r") as f:
                classes = [int(line.split()[0]) for line in f if line.strip()]
            # Classes 2..8 are rarer dialog / modal buttons
            if any(c in {2, 3, 4, 5, 6, 7, 8} for c in classes):
                rare_files.add(fname)

    # Remaining candidate files for validation
    remaining_files = [f for f in image_files if f not in rare_files]
    val_count = max(1, int(len(image_files) * val_ratio))
    val_files = set(remaining_files[:val_count])

    print(f"Preparing dataset splits: {len(image_files) - len(val_files)} train, {len(val_files)} val...")
    print(f"Guaranteed {len(rare_files)} rare modal/dialog images in dataset.")

    for fname in image_files:
        base_name = os.path.splitext(fname)[0]
        src_img = os.path.join(RAW_DIR, fname)
        src_lbl = os.path.join(LABELS_DIR, f"{base_name}.txt")

        is_val = fname in val_files
        dst_img_dir = images_val if is_val else images_train
        dst_lbl_dir = labels_val if is_val else labels_train

        shutil.copy2(src_img, os.path.join(dst_img_dir, fname))
        if os.path.exists(src_lbl):
            shutil.copy2(src_lbl, os.path.join(dst_lbl_dir, f"{base_name}.txt"))
        else:
            open(os.path.join(dst_lbl_dir, f"{base_name}.txt"), "w").close()

    # Oversample rare modal dialog images into train (5x copies with unique names)
    # This guarantees the model learns rare buttons (btn_ok, btn_leave, btn_collect, etc.)
    print(f"Oversampling {len(rare_files)} rare dialog images (5x) into training set...")
    for fname in rare_files:
        base_name = os.path.splitext(fname)[0]
        src_img = os.path.join(RAW_DIR, fname)
        src_lbl = os.path.join(LABELS_DIR, f"{base_name}.txt")
        for copy_idx in range(1, 6):
            copy_img_name = f"{base_name}_copy{copy_idx}.png"
            copy_lbl_name = f"{base_name}_copy{copy_idx}.txt"
            shutil.copy2(src_img, os.path.join(images_train, copy_img_name))
            if os.path.exists(src_lbl):
                shutil.copy2(src_lbl, os.path.join(labels_train, copy_lbl_name))

    print("[SUCCESS] Dataset splits and rare class oversampling completed successfully.")


def train_and_export(epochs: int = 40, batch_size: int = 8, img_size: int = 640):
    """Trains YOLOv8n optimized for Game UI and exports to ONNX."""
    try:
        from ultralytics import YOLO
    except ImportError:
        print("\n[ERROR] 'ultralytics' library not installed!")
        print("Please install it using: pip install ultralytics")
        sys.exit(1)

    yaml_path = os.path.join(DATASET_DIR, "data.yaml")
    os.makedirs(MODELS_DIR, exist_ok=True)

    print("\n" + "=" * 65)
    print("      STARTING UI-OPTIMIZED YOLOV8-NANO TRAINING")
    print("=" * 65)
    print(f"Config: {yaml_path}")
    print(f"Epochs: {epochs} | Batch: {batch_size} | ImgSize: {img_size}")
    print("Augmentations: mosaic=0.0, fliplr=0.0, flipud=0.0 (Preserve UI layout)")
    print("=" * 65)

    # 1. Load nano model pre-trained weights
    model = YOLO("yolov8n.pt")

    # 2. Train model (workers=0 for safe Windows multiprocessing)
    results = model.train(
        data=yaml_path,
        epochs=epochs,
        batch=batch_size,
        imgsz=img_size,
        device="cpu",
        project="runs/detect",
        name="dl_ui_model",
        exist_ok=True,
        workers=0,
        verbose=True,
        # Game UI specific settings:
        mosaic=0.0,   # Never cut/stitch game screens
        fliplr=0.0,   # Never flip horizontally (text/buttons have fixed orientation)
        flipud=0.0,   # Never flip upside down
        degrees=0.0,  # Never rotate UI
    )

    print("\n[TRAINING FINISHED] Exporting to ONNX format...")

    # 3. Export to ONNX
    best_pt = "runs/detect/dl_ui_model/weights/best.pt"
    if not os.path.exists(best_pt):
        # Check alternative nested path
        nested_pt = "runs/detect/runs/detect/dl_ui_model/weights/best.pt"
        if os.path.exists(nested_pt):
            best_pt = nested_pt

    if os.path.exists(best_pt):
        best_model = YOLO(best_pt)
        onnx_path = best_model.export(format="onnx", imgsz=img_size, simplify=True)
        dst_onnx = os.path.join(MODELS_DIR, "dirty_league_yolo.onnx")
        if onnx_path and os.path.exists(onnx_path):
            shutil.copy2(onnx_path, dst_onnx)
            print(f"\n[SUCCESS] Production ONNX model ready at: {dst_onnx}")
            return dst_onnx
    else:
        print(f"[WARNING] {best_pt} not found, checking default export.")
    return None


if __name__ == "__main__":
    prepare_dataset_splits()
    train_and_export(epochs=40, batch_size=8, img_size=640)