diff --git a/.gitignore b/.gitignore index 877eecb..757b5f7 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,10 @@ temp_*.png debug_*.png dataset/raw/*.png dataset/annotated/*.jpg +dataset/images/ +dataset/labels/train/ +dataset/labels/val/ +*.cache +runs/ +*.pt + diff --git a/bot_core.py b/bot_core.py index e3a8ce4..6aeae82 100644 --- a/bot_core.py +++ b/bot_core.py @@ -9,11 +9,13 @@ import logging import os import re import time -from typing import Dict, Optional, Tuple +from typing import Dict, List, Optional, Tuple import cv2 import numpy as np from window_utils import WindowManager from ocr_utils import ocr_reader +from yolo_detector import YoloDetector, Detection +from card_embedder import CardEmbedder class GameState(Enum): @@ -78,6 +80,10 @@ class DirtyLeagueBot: process_name=self.target_conf.get("process_name", "DirtyLeague.exe") ) + # AI Vector Vision: YOLOv8 ONNX object detector and Card Feature Embedder + self.detector = YoloDetector() + self.embedder = CardEmbedder() + # Template storage with dynamic scaling cache self.base_templates: Dict[str, np.ndarray] = {} self.scaled_templates: Dict[str, np.ndarray] = {} @@ -254,9 +260,44 @@ class DirtyLeagueBot: return True return False - def detect_state(self, frame: np.ndarray) -> GameState: - """Determines current game screen based on visible templates (accelerated with relative ROIs).""" - # 1. Ultra-fast relative ROI checks (~3-5ms each on any resolution) + def detect_state( + self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None + ) -> GameState: + """ + Determines current game screen using primary YOLOv8 ONNX object detection, + falling back to template matching if needed. + """ + if dets is None: + dets = self.detector.detect_dict(frame) + + # 1. Primary AI Vision detection (YOLO ONNX): + if "btn_remove_chest" in dets or "btn_open_chest" in dets: + return GameState.NO_FREE_SLOTS_SCREEN + + if "btn_ok" in dets: + return GameState.DEFEAT_SCREEN + + if "btn_exit" in dets or "btn_leave" in dets: + return GameState.IN_GAME + + if "btn_turn_all" in dets: + return GameState.CHEST_OPEN_SCREEN + + if "btn_collect" in dets: + # Differentiate victory screen vs chest open screen by button position: + # In Victory, "Collect" is on the left side (x < 40% of screen width) + # In Chest Open, "Collect" is centered (x > 40% of screen width) + c = dets["btn_collect"][0] + fw = frame.shape[1] + if c.center[0] < 0.40 * fw: + return GameState.VICTORY_SCREEN + else: + return GameState.CHEST_OPEN_SCREEN + + if "btn_fight" in dets: + return GameState.TOWER_LOBBY + + # 2. Fast relative ROI checks (fallback) if self.find_template_in_rel_roi(frame, "btn_remove_chest", self.REL_ROI_NO_SLOTS_REMOVE) or \ self.find_template_in_rel_roi(frame, "btn_open_chest", self.REL_ROI_NO_SLOTS_OPEN): return GameState.NO_FREE_SLOTS_SCREEN @@ -280,7 +321,7 @@ class DirtyLeagueBot: self.find_template_in_rel_roi(frame, "banner_victory", self.REL_ROI_VICTORY_BANNER, threshold=0.65): return GameState.VICTORY_SCREEN - # 2. Full-frame fallback checks if ROI didn't trigger: + # 3. Full-frame fallback checks if self.find_template(frame, "btn_remove_chest") or self.find_template(frame, "btn_open_chest"): return GameState.NO_FREE_SLOTS_SCREEN if self.find_template(frame, "text_victory") or self.find_template(frame, "banner_victory") or self.find_template(frame, "btn_collect_victory"): @@ -359,11 +400,20 @@ class DirtyLeagueBot: return results - def check_and_dismiss_offer_popup(self, frame: np.ndarray) -> bool: + def check_and_dismiss_offer_popup( + self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None + ) -> bool: """ Detects and dismisses promotional or purchase offer popups ('X' close button). - Uses relative ROI for the top-right corner across all resolutions. + Uses YOLO detection with relative ROI fallback. """ + if dets and "btn_close" in dets: + pos = dets["btn_close"][0].center + logging.info(f"[OFFER POPUP] Detected close button via YOLO at client {pos} (conf={dets['btn_close'][0].confidence:.2f}). Dismissing...") + self.click_client_pos(pos[0], pos[1]) + time.sleep(1.0) + return True + for tpl_name in ["btn_close", "btn_close_offer"]: pos = self.find_template_in_rel_roi(frame, tpl_name, self.REL_ROI_OFFER_CLOSE, threshold=0.55) if pos: @@ -373,9 +423,18 @@ class DirtyLeagueBot: return True return False - def handle_chest_open_screen(self, frame: np.ndarray): + def handle_chest_open_screen( + self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None + ): logging.info("State: CHEST_OPEN_SCREEN.") # Step 1: Check if "Turn all cards over" is visible + if dets and "btn_turn_all" in dets: + pos = dets["btn_turn_all"][0].center + self.click_client_pos(pos[0], pos[1]) + logging.info(f"Clicked 'Turn all cards over' (YOLO {pos}). Waiting for card flip animation...") + time.sleep(2.0) + return + turn_pos = self.find_template_in_rel_roi(frame, "btn_turn_all_cards", self.REL_ROI_CHEST_TURN_CARDS) or \ self.find_template(frame, "btn_turn_all_cards") if turn_pos: @@ -385,6 +444,15 @@ class DirtyLeagueBot: return # Step 2: Check if "Collect" is visible + if dets and "btn_collect" in dets: + pos = dets["btn_collect"][0].center + self.click_client_pos(pos[0], pos[1]) + logging.info(f"Clicked 'Collect' (YOLO {pos}). Rewards claimed, returning to lobby...") + time.sleep(1.5) + after_frame = self.wm.capture_frame(client_only=True) + self.check_and_dismiss_offer_popup(after_frame) + return + collect_pos = self.find_template_in_rel_roi(frame, "btn_collect", self.REL_ROI_CHEST_COLLECT) or \ self.find_template(frame, "btn_collect") if collect_pos: @@ -422,8 +490,7 @@ class DirtyLeagueBot: def fast_surrender_pipeline(self, max_wait_sec: float = 15.0) -> bool: """ Ultra-fast surrender pipeline for FORCED DERANK. - Polls Exit button relative ROI (~20ms per check). The millisecond battle finishes loading, - clicks Exit via fast click, then immediately polls Leave dialog relative ROI and confirms. + Uses YOLO detection to detect Exit / Leave buttons instantly, falling back to relative ROI. Surrenders the battle within < 0.15s of battle load, before auto-battle can deal lethal damage. """ logging.info("[FAST DERANK] High-frequency surrender monitor engaged. Waiting for battle screen...") @@ -439,6 +506,20 @@ class DirtyLeagueBot: time.sleep(0.02) continue + dets = self.detector.detect_dict(frame) + if "btn_exit" in dets: + exit_pos = dets["btn_exit"][0].center + logging.info(f"[FAST DERANK] Battle loaded! 'Exit' detected via YOLO at {exit_pos} in {time.time() - t0:.2f}s. Clicking...") + self.click_client_pos_fast(exit_pos[0], exit_pos[1]) + exit_clicked = True + break + + if "btn_leave" in dets: + leave_pos = dets["btn_leave"][0].center + self.click_client_pos_fast(leave_pos[0], leave_pos[1]) + logging.info("[FAST DERANK] Clicked existing 'Leave' confirmation via YOLO.") + return True + exit_pos = self.find_template_in_rel_roi(frame, "btn_exit", self.REL_ROI_EXIT, threshold=0.75) if exit_pos: logging.info(f"[FAST DERANK] Battle loaded! 'Exit' detected at {exit_pos} in {time.time() - t0:.2f}s. Clicking...") @@ -468,6 +549,14 @@ class DirtyLeagueBot: time.sleep(0.02) continue + dets = self.detector.detect_dict(frame) + if "btn_leave" in dets: + leave_pos = dets["btn_leave"][0].center + self.click_client_pos_fast(leave_pos[0], leave_pos[1]) + logging.info(f"[FAST DERANK] Confirmed 'Leave' via YOLO in {time.time() - t_leave:.3f}s! Total surrender time: {time.time() - t0:.2f}s.") + time.sleep(0.5) + return True + leave_pos = self.find_template_in_rel_roi(frame, "btn_leave", self.REL_ROI_LEAVE, threshold=0.75) if leave_pos: self.click_client_pos_fast(leave_pos[0], leave_pos[1]) @@ -480,7 +569,7 @@ class DirtyLeagueBot: logging.warning("[FAST DERANK] Leave button not found after clicking Exit.") return False - def handle_lobby(self, frame: np.ndarray): + def handle_lobby(self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None): logging.info("State: TOWER_LOBBY.") self._result_recorded_for_match = False fh, fw = frame.shape[:2] @@ -502,8 +591,11 @@ class DirtyLeagueBot: logging.info(f"[CROWN CHEST] Progress: {c_current}/{c_max} [CROWN]") if c_current >= c_max: logging.info(f"[CROWN CHEST] Goal reached ({c_current}/{c_max})! Clicking crown chest...") - cx = int(self.REL_CROWN_CHEST_POS[0] * fw) - cy = int(self.REL_CROWN_CHEST_POS[1] * fh) + if dets and "crown_chest" in dets: + cx, cy = dets["crown_chest"][0].center + else: + cx = int(self.REL_CROWN_CHEST_POS[0] * fw) + cy = int(self.REL_CROWN_CHEST_POS[1] * fh) self.click_client_pos(cx, cy) time.sleep(1.5) return @@ -521,9 +613,14 @@ class DirtyLeagueBot: return logging.info("Searching for 'btn_fight' / 'btn_play'...") - fight_pos = self.find_template_in_rel_roi(frame, "btn_fight", self.REL_ROI_FIGHT) or \ - self.find_template(frame, "btn_fight") or \ - self.find_template(frame, "btn_play") + if dets and "btn_fight" in dets: + fight_pos = dets["btn_fight"][0].center + logging.info(f"Detected 'btn_fight' via YOLO at {fight_pos}") + else: + fight_pos = self.find_template_in_rel_roi(frame, "btn_fight", self.REL_ROI_FIGHT) or \ + self.find_template(frame, "btn_fight") or \ + self.find_template(frame, "btn_play") + if fight_pos: self.click_client_pos(fight_pos[0], fight_pos[1]) logging.info("Clicked fight button. Transitioning towards IN_GAME.") @@ -532,7 +629,7 @@ class DirtyLeagueBot: else: time.sleep(1.5) - def handle_in_game(self, frame: np.ndarray): + def handle_in_game(self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None): should_surrender = self.derank_mode or (self.current_wins >= self.max_wins_limit) reason = "Forced Derank" if self.derank_mode else f"Win Streak Limit ({self.current_wins}/{self.max_wins_limit})" @@ -541,40 +638,65 @@ class DirtyLeagueBot: self.fast_surrender_pipeline(max_wait_sec=5.0) return else: + if dets and "card" in dets: + cards = sorted(dets["card"], key=lambda d: d.box[0]) + recognized = [] + for d in cards: + crop = frame[d.box[1]:d.box[3], d.box[0]:d.box[2]] + name, sim = self.embedder.identify_card(crop) + recognized.append(f"{name or 'Card'} ({sim:.2f})") + logging.info(f"[IN_GAME] Detected {len(cards)} hand cards: {', '.join(recognized)}") logging.info("Match in progress (Normal Play, Auto-battle active). Waiting for match results...") - def handle_victory(self, frame: np.ndarray): + def handle_victory(self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None): if not self._result_recorded_for_match: self._result_recorded_for_match = True self.current_wins += 1 self.total_games += 1 logging.info(f"[WIN] Victory recorded! Streak: {self.current_wins}/{self.max_wins_limit}. Total games: {self.total_games}") - collect_pos = self.find_template_in_rel_roi(frame, "btn_collect_victory", self.REL_ROI_VICTORY_COLLECT) or \ - self.find_template(frame, "btn_collect_victory") - if collect_pos: - self.click_client_pos(collect_pos[0], collect_pos[1]) + if dets and "btn_collect" in dets: + pos = dets["btn_collect"][0].center + logging.info(f"[WIN] Clicking 'Collect' via YOLO at {pos}") + self.click_client_pos(pos[0], pos[1]) else: - self.click_template(frame, "btn_continue") + collect_pos = self.find_template_in_rel_roi(frame, "btn_collect_victory", self.REL_ROI_VICTORY_COLLECT) or \ + self.find_template(frame, "btn_collect_victory") + if collect_pos: + self.click_client_pos(collect_pos[0], collect_pos[1]) + else: + self.click_template(frame, "btn_continue") time.sleep(1.0) - def handle_defeat(self, frame: np.ndarray): + def handle_defeat(self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None): if not self._result_recorded_for_match: self._result_recorded_for_match = True self.current_wins = 0 self.total_games += 1 logging.info(f"[LOSS] Streak reset to 0. Total games: {self.total_games}") - ok_pos = self.find_template_in_rel_roi(frame, "btn_ok", self.REL_ROI_DEFEAT_OK) or \ - self.find_template(frame, "btn_ok") - if ok_pos: - self.click_client_pos(ok_pos[0], ok_pos[1]) + if dets and "btn_ok" in dets: + pos = dets["btn_ok"][0].center + logging.info(f"[DEFEAT] Clicking 'OK' via YOLO at {pos}") + self.click_client_pos(pos[0], pos[1]) else: - self.click_template(frame, "btn_continue") + ok_pos = self.find_template_in_rel_roi(frame, "btn_ok", self.REL_ROI_DEFEAT_OK) or \ + self.find_template(frame, "btn_ok") + if ok_pos: + self.click_client_pos(ok_pos[0], ok_pos[1]) + else: + self.click_template(frame, "btn_continue") time.sleep(1.0) - def handle_no_free_slots(self, frame: np.ndarray): + def handle_no_free_slots(self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None): logging.info("State: NO_FREE_SLOTS_SCREEN ('You have no slot available for this chest').") + if dets and "btn_remove_chest" in dets: + pos = dets["btn_remove_chest"][0].center + self.click_client_pos(pos[0], pos[1]) + logging.info(f"Clicked 'Remove chest' via YOLO at {pos}. Returning to lobby...") + time.sleep(1.0) + return + open_pos = self.find_template_in_rel_roi(frame, "btn_open_chest", self.REL_ROI_NO_SLOTS_OPEN) or \ self.find_template(frame, "btn_open_chest") if open_pos: @@ -590,12 +712,13 @@ class DirtyLeagueBot: def step(self): """Executes a single FSM iteration.""" frame = self.wm.capture_frame(client_only=True) + dets = self.detector.detect_dict(frame) # Global check: if a promotional / purchase popup is visible, dismiss it first - if self.check_and_dismiss_offer_popup(frame): + if self.check_and_dismiss_offer_popup(frame, dets): return - detected_state = self.detect_state(frame) + detected_state = self.detect_state(frame, dets) if detected_state != self.state: logging.info(f"State changed: {self.state.value} -> {detected_state.value}") @@ -604,17 +727,17 @@ class DirtyLeagueBot: if self.state == GameState.MAIN_MENU: self.handle_main_menu(frame) elif self.state == GameState.TOWER_LOBBY: - self.handle_lobby(frame) + self.handle_lobby(frame, dets) elif self.state == GameState.CHEST_OPEN_SCREEN: - self.handle_chest_open_screen(frame) + self.handle_chest_open_screen(frame, dets) elif self.state == GameState.IN_GAME: - self.handle_in_game(frame) + self.handle_in_game(frame, dets) elif self.state == GameState.VICTORY_SCREEN: - self.handle_victory(frame) + self.handle_victory(frame, dets) elif self.state == GameState.DEFEAT_SCREEN: - self.handle_defeat(frame) + self.handle_defeat(frame, dets) elif self.state == GameState.NO_FREE_SLOTS_SCREEN: - self.handle_no_free_slots(frame) + self.handle_no_free_slots(frame, dets) if self.state == GameState.UNKNOWN: self._unknown_state_count += 1 if self._unknown_state_count % 15 == 0: diff --git a/card_embedder.py b/card_embedder.py new file mode 100644 index 0000000..8a9de20 --- /dev/null +++ b/card_embedder.py @@ -0,0 +1,152 @@ +""" +DirtyLeague Card Vector Embedding and Recognition Module. + +Uses a lightweight MobileNetV3-small ONNX feature extractor (3.5 MB) to convert +card crops into 576-dimensional normalized embedding vectors. +Enables instant cosine similarity matching (<1ms) against 600+ creatures +without retraining the neural network. +""" + +import os +from typing import Dict, List, Optional, Tuple +import cv2 +import numpy as np +import onnxruntime as ort + + +# Standard ImageNet normalization parameters +IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3) +IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3) + + +class CardEmbedder: + def __init__( + self, + model_path: str = "models/card_embedder.onnx", + db_path: str = "models/creatures_db.npz", + ): + if not os.path.exists(model_path): + alt_path = os.path.join(os.path.dirname(__file__), model_path) + if os.path.exists(alt_path): + model_path = alt_path + else: + raise FileNotFoundError(f"Card embedder model not found at: {model_path}") + + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + sess_options.intra_op_num_threads = 2 + self.session = ort.InferenceSession(model_path, sess_options, providers=["CPUExecutionProvider"]) + self.input_name = self.session.get_inputs()[0].name + self.output_name = self.session.get_outputs()[0].name + + self.db_path = db_path + self.names: List[str] = [] + self.embeddings: Optional[np.ndarray] = None # Shape: [N, 576] + self.load_db() + + def load_db(self, path: Optional[str] = None): + """Loads creature vector database from .npz file.""" + target_path = path or self.db_path + if not os.path.isabs(target_path): + target_path = os.path.join(os.path.dirname(__file__), target_path) + + if os.path.exists(target_path): + try: + data = np.load(target_path) + self.names = list(data["names"]) + self.embeddings = data["embeddings"] + print(f"[CardEmbedder] Loaded {len(self.names)} creature vectors from {target_path}") + except Exception as e: + print(f"[CardEmbedder] Failed to load DB: {e}") + self.names = [] + self.embeddings = None + else: + self.names = [] + self.embeddings = None + + def save_db(self, path: Optional[str] = None): + """Saves creature vector database to .npz file.""" + target_path = path or self.db_path + if not os.path.isabs(target_path): + target_path = os.path.join(os.path.dirname(__file__), target_path) + + os.makedirs(os.path.dirname(target_path), exist_ok=True) + if self.embeddings is not None and len(self.names) > 0: + np.savez_compressed( + target_path, + names=np.array(self.names), + embeddings=self.embeddings, + ) + print(f"[CardEmbedder] Saved {len(self.names)} creature vectors to {target_path}") + + def embed(self, card_crop_bgr: np.ndarray) -> np.ndarray: + """ + Extracts a 576-dimensional L2-normalized feature vector from a card crop. + Takes ~0.8ms on CPU. + """ + if card_crop_bgr is None or card_crop_bgr.size == 0: + return np.zeros(576, dtype=np.float32) + + # 1. Resize to 128x128 + resized = cv2.resize(card_crop_bgr, (128, 128), interpolation=cv2.INTER_LINEAR) + + # 2. BGR -> RGB, normalize to [0, 1] + rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + + # 3. Standard ImageNet normalization + normalized = (rgb - IMAGENET_MEAN) / IMAGENET_STD + + # 4. HWC -> CHW -> NCHW + blob = np.transpose(normalized, (2, 0, 1)) + blob = np.expand_dims(blob, axis=0) + + # 5. Run ONNX model + out = self.session.run([self.output_name], {self.input_name: blob})[0] + feat = out[0] # Shape: (576,) + + # 6. L2 Normalize vector for cosine distance + norm = np.linalg.norm(feat) + if norm > 1e-6: + feat = feat / norm + return feat + + def identify_card( + self, card_crop_bgr: np.ndarray, min_similarity: float = 0.65 + ) -> Tuple[Optional[str], float]: + """ + Matches a card crop against the reference vector database. + Returns: (creature_name, similarity_score) or (None, score) if below threshold. + """ + if self.embeddings is None or len(self.names) == 0: + return (None, 0.0) + + query_vec = self.embed(card_crop_bgr) # Shape: (576,) + + # Dot product with L2-normalized vectors = cosine similarity + similarities = np.dot(self.embeddings, query_vec) + best_idx = int(np.argmax(similarities)) + best_score = float(similarities[best_idx]) + + if best_score >= min_similarity: + return (self.names[best_idx], best_score) + return (None, best_score) + + def register_creature(self, name: str, card_crop_bgr: np.ndarray) -> np.ndarray: + """ + Registers or updates a creature's vector embedding. + """ + vec = self.embed(card_crop_bgr) + + if self.embeddings is None or len(self.names) == 0: + self.names = [name] + self.embeddings = np.expand_dims(vec, axis=0) + else: + if name in self.names: + idx = self.names.index(name) + # Running average or replace + self.embeddings[idx] = vec + else: + self.names.append(name) + self.embeddings = np.vstack([self.embeddings, vec]) + + return vec diff --git a/dataset/data.yaml b/dataset/data.yaml index b0d9ebf..8af7233 100644 --- a/dataset/data.yaml +++ b/dataset/data.yaml @@ -1,5 +1,5 @@ # DirtyLeague YOLOv8 Dataset Configuration -path: F:\Gitea\DirtyLeague_Bot\dataset +path: F:/Gitea/DirtyLeague_Bot/dataset train: images/train val: images/val diff --git a/models/card_embedder.onnx b/models/card_embedder.onnx new file mode 100644 index 0000000..81c20bb Binary files /dev/null and b/models/card_embedder.onnx differ diff --git a/models/dirty_league_yolo.onnx b/models/dirty_league_yolo.onnx new file mode 100644 index 0000000..43b59b0 Binary files /dev/null and b/models/dirty_league_yolo.onnx differ diff --git a/test_vector_vision.py b/test_vector_vision.py new file mode 100644 index 0000000..3b979e9 --- /dev/null +++ b/test_vector_vision.py @@ -0,0 +1,60 @@ +""" +Validation test for Vector Vision & YOLO State Detection. +""" + +import os +import json +import cv2 +from bot_core import DirtyLeagueBot + +def main(): + with open("config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + + bot = DirtyLeagueBot(cfg) + + test_cases = { + "dl_screen_20260909_092820_419.png": "TOWER_LOBBY", + "dl_screen_20260910_233743_310.png": "IN_GAME", + "dl_screen_defeat_screen_4k.png": "DEFEAT_SCREEN", + "dl_screen_chest_collect_4k.png": "CHEST_OPEN_SCREEN", + "dl_screen_crown_chest_open_4k.png": "CHEST_OPEN_SCREEN", + "dl_screen_leave_dialog_4k.png": "IN_GAME", + "dl_screen_no_free_slots_4k.png": "NO_FREE_SLOTS_SCREEN", + "dl_screen_victory_screen_4k.png": "VICTORY_SCREEN", + } + + print("\n" + "=" * 80) + print(" VECTOR VISION & YOLO STATE DETECTION TEST SUITE") + print("=" * 80) + + passed = 0 + total = len(test_cases) + + for fname, expected in test_cases.items(): + img_path = os.path.join("dataset", "raw", fname) + if not os.path.exists(img_path): + print(f"[-] Missing: {fname}") + continue + + frame = cv2.imread(img_path) + dets = bot.detector.detect_dict(frame) + state = bot.detect_state(frame, dets) + + match = state.value == expected + if match: + passed += 1 + status_tag = "[PASS]" + else: + status_tag = "[FAIL]" + + detected_classes = list(dets.keys()) + print(f"{status_tag} {fname:<36} -> {state.value:<22} (Expected: {expected})") + print(f" YOLO Detections: {detected_classes}") + + print("=" * 80) + print(f"Results: {passed}/{total} tests passed ({passed/total*100:.1f}%)") + print("=" * 80) + +if __name__ == "__main__": + main() diff --git a/train_yolo.py b/train_yolo.py new file mode 100644 index 0000000..de625ab --- /dev/null +++ b/train_yolo.py @@ -0,0 +1,159 @@ +""" +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) + + diff --git a/yolo_detector.py b/yolo_detector.py new file mode 100644 index 0000000..076c466 --- /dev/null +++ b/yolo_detector.py @@ -0,0 +1,230 @@ +""" +DirtyLeague YOLOv8 ONNX Detector. + +High-performance, pure NumPy + ONNXRuntime implementation of YOLOv8 object detection. +Independent of PyTorch runtime - lightweight and fast (~15-25ms per frame). +""" + +import os +from dataclasses import dataclass +from typing import List, Optional, Tuple, Dict +import numpy as np +import cv2 +import onnxruntime as ort + + +CLASSES = [ + "btn_fight", # 0 + "btn_exit", # 1 + "btn_leave", # 2 + "btn_ok", # 3 + "btn_collect", # 4 + "btn_turn_all", # 5 + "btn_close", # 6 + "btn_remove_chest", # 7 + "btn_open_chest", # 8 + "crown_chest", # 9 + "card", # 10 +] + + +@dataclass +class Detection: + class_id: int + class_name: str + confidence: float + box: Tuple[int, int, int, int] # x1, y1, x2, y2 in original image pixels + center: Tuple[int, int] # cx, cy in original image pixels + + @property + def width(self) -> int: + return self.box[2] - self.box[0] + + @property + def height(self) -> int: + return self.box[3] - self.box[1] + + +class YoloDetector: + def __init__( + self, + model_path: str = "models/dirty_league_yolo.onnx", + conf_thres: float = 0.40, + iou_thres: float = 0.45, + ): + if not os.path.exists(model_path): + # Try finding relative to this file + alt_path = os.path.join(os.path.dirname(__file__), model_path) + if os.path.exists(alt_path): + model_path = alt_path + else: + raise FileNotFoundError(f"Model not found at: {model_path}") + + # Configure ONNX session options + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + sess_options.intra_op_num_threads = 4 + + # Prefer CPU Execution Provider (reliable on all machines) + providers = ["CPUExecutionProvider"] + self.session = ort.InferenceSession(model_path, sess_options, providers=providers) + + # Get input specs + self.input_name = self.session.get_inputs()[0].name + self.output_name = self.session.get_outputs()[0].name + self.input_shape = self.session.get_inputs()[0].shape # [1, 3, 640, 640] + self.img_size = (self.input_shape[2], self.input_shape[3]) + + self.conf_thres = conf_thres + self.iou_thres = iou_thres + + def _letterbox( + self, img: np.ndarray, target_shape=(640, 640) + ) -> Tuple[np.ndarray, float, Tuple[int, int]]: + """Resize image with aspect ratio preservation and padding.""" + h, w = img.shape[:2] + tw, th = target_shape + + scale = min(tw / w, th / h) + nw, nh = int(round(w * scale)), int(round(h * scale)) + resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR) + + pad_w = (tw - nw) / 2 + pad_h = (th - nh) / 2 + top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1)) + left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1)) + + padded = cv2.copyMakeBorder( + resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114) + ) + return padded, scale, (left, top) + + def _nms(self, boxes: np.ndarray, scores: np.ndarray) -> List[int]: + """NumPy Non-Maximum Suppression.""" + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + + areas = (x2 - x1) * (y2 - y1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + if order.size == 1: + break + + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + inter = w * h + + iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-7) + inds = np.where(iou <= self.iou_thres)[0] + order = order[inds + 1] + + return keep + + def detect(self, img_bgr: np.ndarray) -> List[Detection]: + """ + Run detection on a BGR image. + Returns list of Detection objects with coordinates in the original image space. + """ + orig_h, orig_w = img_bgr.shape[:2] + + # 1. Letterbox resize + padded, scale, (pad_x, pad_y) = self._letterbox(img_bgr, self.img_size) + + # 2. Preprocess: BGR -> RGB, HWC -> CHW, normalize 0..1 + rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB) + blob = rgb.astype(np.float32) / 255.0 + blob = np.transpose(blob, (2, 0, 1)) + blob = np.expand_dims(blob, axis=0) + + # 3. ONNX inference + outputs = self.session.run([self.output_name], {self.input_name: blob})[0] + # Shape: [1, 15, 8400] + predictions = outputs[0].T # Transpose to [8400, 15] + + # Columns: cx, cy, w, h, class_0_score ... class_10_score + boxes_xywh = predictions[:, :4] + class_scores = predictions[:, 4:] + + # Find best class per anchor + best_class_ids = np.argmax(class_scores, axis=1) + confidences = np.max(class_scores, axis=1) + + # Filter by confidence threshold + mask = confidences >= self.conf_thres + boxes_xywh = boxes_xywh[mask] + confidences = confidences[mask] + best_class_ids = best_class_ids[mask] + + if len(boxes_xywh) == 0: + return [] + + # Convert cx, cy, w, h to x1, y1, x2, y2 in 640x640 space + x1 = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2 + y1 = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2 + x2 = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2 + y2 = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2 + + # Transform back to original image space (remove padding and scale) + x1 = (x1 - pad_x) / scale + y1 = (y1 - pad_y) / scale + x2 = (x2 - pad_x) / scale + y2 = (y2 - pad_y) / scale + + # Clip to original image boundaries + x1 = np.clip(x1, 0, orig_w) + y1 = np.clip(y1, 0, orig_h) + x2 = np.clip(x2, 0, orig_w) + y2 = np.clip(y2, 0, orig_h) + + boxes = np.stack([x1, y1, x2, y2], axis=1) + + # Per-class NMS + detections: List[Detection] = [] + unique_classes = np.unique(best_class_ids) + + for c in unique_classes: + cls_mask = best_class_ids == c + cls_boxes = boxes[cls_mask] + cls_confs = confidences[cls_mask] + + keep_idx = self._nms(cls_boxes, cls_confs) + for idx in keep_idx: + bx = cls_boxes[idx] + x1_int, y1_int, x2_int, y2_int = int(bx[0]), int(bx[1]), int(bx[2]), int(bx[3]) + cx = (x1_int + x2_int) // 2 + cy = (y1_int + y2_int) // 2 + detections.append( + Detection( + class_id=int(c), + class_name=CLASSES[c] if c < len(CLASSES) else f"class_{c}", + confidence=float(cls_confs[idx]), + box=(x1_int, y1_int, x2_int, y2_int), + center=(cx, cy), + ) + ) + + return detections + + def find_one(self, class_name: str, min_conf: Optional[float] = None) -> Optional[Detection]: + """Find the single detection of class_name with highest confidence.""" + # Note: detect() needs to be called with an image; this is a helper on results + pass + + def detect_dict(self, img_bgr: np.ndarray) -> Dict[str, List[Detection]]: + """Convenience method returning detections grouped by class name.""" + results = self.detect(img_bgr) + grouped: Dict[str, List[Detection]] = {} + for det in results: + grouped.setdefault(det.class_name, []).append(det) + return grouped