""" 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)