volandsz/DirtyLeague_Bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(vision): integrate YOLOv8 ONNX object detector and MobileNetV3 card embedder

Browse Source
This commit is contained in:
VolandSZ
2026-09-11 00:40:10 +03:00
parent 35b0e53413
commit 3a40c2003f
9 changed files with 770 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
.gitignore vendored
View File

@ -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

199
bot_core.py
View File

@ -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:

152
card_embedder.py Normal file
View File

@ -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

2
dataset/data.yaml
View File

@ -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

BIN
models/card_embedder.onnx Normal file
View File

Binary file not shown.

BIN
models/dirty_league_yolo.onnx Normal file
View File

Binary file not shown.

60
test_vector_vision.py Normal file
View File

@ -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()

159
train_yolo.py Normal file
View File

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

230
yolo_detector.py Normal file
View File

@ -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