Compare commits

4 Commits

83 changed files with 1515 additions and 38 deletions

10
.gitignore vendored
View File

@ -19,3 +19,13 @@ logs/
*.log *.log
temp_*.png temp_*.png
debug_*.png debug_*.png
debug_cards/
dataset/raw/*.png
dataset/annotated/*.jpg
dataset/images/
dataset/labels/train/
dataset/labels/val/
*.cache
runs/
*.pt

259
auto_label.py Normal file
View File

@ -0,0 +1,259 @@
"""
Auto-Labeling Assistant for DirtyLeague Dataset.
Uses calibrated ROIs and templates to automatically generate
YOLO-format annotations (.txt) for the collected raw screenshots in seconds.
Outputs:
- dataset/labels/<image_name>.txt (normalized YOLO bbox: class_id cx cy w h)
- dataset/annotated/<image_name>.jpg (visual verification with labeled boxes)
- dataset/data.yaml (dataset configuration for YOLO training)
"""
import os
import sys
import cv2
import numpy as np
# Class definitions for DirtyLeague UI detection
CLASSES = [
"btn_fight", # 0: Yellow/Gold fight button in lobby
"btn_exit", # 1: White battle exit button (bottom bar)
"btn_leave", # 2: Red confirm surrender button
"btn_ok", # 3: Defeat screen OK button
"btn_collect", # 4: Victory / Chest collect button
"btn_turn_all", # 5: Chest turn cards over button
"btn_close", # 6: Offer / modal close 'X' button
"btn_remove_chest", # 7: Discard extra chest button
"btn_open_chest", # 8: Open chest button on no-slots screen
"crown_chest", # 9: 5/5 Crown chest banner in lobby
"card", # 10: Generic creature card / hero slot
]
CLASS_MAP = {name: i for i, name in enumerate(CLASSES)}
# Template mappings to class names
TEMPLATE_TO_CLASS = {
"btn_fight": "btn_fight",
"btn_exit": "btn_exit",
"btn_leave": "btn_leave",
"btn_ok": "btn_ok",
"btn_collect": "btn_collect",
"btn_collect_victory": "btn_collect",
"btn_turn_all_cards": "btn_turn_all",
"btn_close": "btn_close",
"btn_close_offer": "btn_close",
"btn_remove_chest": "btn_remove_chest",
"btn_open_chest": "btn_open_chest",
}
# Calibrated Relative ROIs (rx1, ry1, rx2, ry2) for lightning-fast localized search (<1ms)
TEMPLATE_ROIS = {
"btn_fight": [(0.75, 0.35, 0.95, 0.50)],
"btn_exit": [(0.40, 0.90, 0.52, 1.00)],
"btn_leave": [(0.48, 0.54, 0.64, 0.70)],
"btn_ok": [(0.20, 0.70, 0.40, 0.85)],
"btn_collect": [(0.38, 0.63, 0.60, 0.78)],
"btn_collect_victory": [(0.18, 0.73, 0.33, 0.87)],
"btn_turn_all_cards": [(0.38, 0.63, 0.60, 0.78)],
"btn_close": [(0.85, 0.00, 1.00, 0.16)],
"btn_close_offer": [(0.85, 0.00, 1.00, 0.16)],
"btn_remove_chest": [(0.33, 0.80, 0.52, 0.92)],
"btn_open_chest": [(0.46, 0.80, 0.63, 0.92)],
}
# Color palette for visual annotations
COLORS = [
(0, 215, 255), # btn_fight (gold)
(255, 255, 255), # btn_exit (white)
(0, 0, 255), # btn_leave (red)
(255, 100, 0), # btn_ok (blue)
(0, 200, 0), # btn_collect (green)
(255, 0, 255), # btn_turn_all (magenta)
(0, 165, 255), # btn_close (orange)
(128, 0, 128), # btn_remove_chest (purple)
(200, 200, 0), # btn_open_chest (cyan)
(50, 205, 50), # crown_chest (lime)
(255, 191, 0), # card (deep sky blue)
]
def load_templates(assets_dir: str):
templates = {}
for tpl_name in TEMPLATE_TO_CLASS.keys():
path = os.path.join(assets_dir, f"{tpl_name}.png")
if os.path.exists(path):
img = cv2.imread(path)
if img is not None:
templates[tpl_name] = img
return templates
def detect_boxes(image: np.ndarray, templates: dict, base_height: int = 2050):
fh, fw = image.shape[:2]
scale = fh / float(base_height)
boxes = [] # list of (class_id, x1, y1, x2, y2, conf)
# 1. Fast localized Template Matching via ROIs
for tpl_name, base_tpl in templates.items():
class_name = TEMPLATE_TO_CLASS[tpl_name]
class_id = CLASS_MAP[class_name]
interp = cv2.INTER_AREA if scale < 1.0 else cv2.INTER_CUBIC
tpl = cv2.resize(base_tpl, (0, 0), fx=scale, fy=scale, interpolation=interp)
th, tw = tpl.shape[:2]
rois = TEMPLATE_ROIS.get(tpl_name, [(0.0, 0.0, 1.0, 1.0)])
threshold = 0.55 if "close" in tpl_name else 0.70
for rx1, ry1, rx2, ry2 in rois:
crop_x1 = max(0, min(fw, int(rx1 * fw)))
crop_y1 = max(0, min(fh, int(ry1 * fh)))
crop_x2 = max(0, min(fw, int(rx2 * fw)))
crop_y2 = max(0, min(fh, int(ry2 * fh)))
if (crop_x2 - crop_x1) < tw or (crop_y2 - crop_y1) < th:
continue
roi_crop = image[crop_y1:crop_y2, crop_x1:crop_x2]
res = cv2.matchTemplate(roi_crop, tpl, cv2.TM_CCOEFF_NORMED)
min_v, max_v, min_l, max_l = cv2.minMaxLoc(res)
if max_v >= threshold:
x1 = crop_x1 + max_l[0]
y1 = crop_y1 + max_l[1]
x2 = x1 + tw
y2 = y1 + th
boxes.append((class_id, x1, y1, x2, y2, float(max_v)))
# 2. Geometric Anchor for Crown Chest (in lobby screen if btn_fight is detected)
has_fight = any(b[0] == CLASS_MAP["btn_fight"] for b in boxes)
if has_fight:
cx = int(0.2279 * fw)
cy = int(0.6146 * fh)
cw = int(0.12 * fw)
ch = int(0.14 * fh)
x1, y1 = cx - cw // 2, cy - ch // 2
x2, y2 = cx + cw // 2, cy + ch // 2
boxes.append((CLASS_MAP["crown_chest"], x1, y1, x2, y2, 0.95))
# 3. Detect cards in Battle Screen (if btn_exit is detected)
has_exit = any(b[0] == CLASS_MAP["btn_exit"] for b in boxes)
if has_exit:
# Detect card slots along the bottom player hand / deck bar
# In battle, player deck cards are anchored around y = 0.78..0.92, spaced across bottom
card_w = int(140 * scale)
card_h = int(190 * scale)
card_y = int(0.85 * fh)
card_xs = [int(x_ratio * fw) for x_ratio in [0.22, 0.32, 0.68, 0.78]]
for cx in card_xs:
x1 = cx - card_w // 2
y1 = card_y - card_h // 2
x2 = cx + card_w // 2
y2 = card_y + card_h // 2
boxes.append((CLASS_MAP["card"], x1, y1, x2, y2, 0.85))
return boxes
def write_yolo_labels(boxes, image_shape, out_txt_path):
fh, fw = image_shape[:2]
lines = []
for class_id, x1, y1, x2, y2, _ in boxes:
cx = ((x1 + x2) / 2.0) / fw
cy = ((y1 + y2) / 2.0) / fh
w = (x2 - x1) / float(fw)
h = (y2 - y1) / float(fh)
lines.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
with open(out_txt_path, "w", encoding="utf-8") as f:
f.writelines(lines)
def draw_annotations(image, boxes):
annotated = image.copy()
for class_id, x1, y1, x2, y2, conf in boxes:
color = COLORS[class_id % len(COLORS)]
cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 3)
label = f"{CLASSES[class_id]} {conf:.2f}"
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
cv2.rectangle(annotated, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, -1)
cv2.putText(annotated, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
return annotated
def main():
root_dir = os.path.dirname(__file__)
raw_dir = os.path.join(root_dir, "dataset", "raw")
labels_dir = os.path.join(root_dir, "dataset", "labels")
annotated_dir = os.path.join(root_dir, "dataset", "annotated")
assets_dir = os.path.join(root_dir, "assets")
os.makedirs(labels_dir, exist_ok=True)
os.makedirs(annotated_dir, exist_ok=True)
templates = load_templates(assets_dir)
print(f"Loaded {len(templates)} templates for ROI-accelerated pre-labeling.")
image_files = sorted([f for f in os.listdir(raw_dir) if f.endswith(".png")])
print(f"Found {len(image_files)} raw images in '{raw_dir}' to process.")
total_boxes = 0
annotated_images_count = 0
for i, fname in enumerate(image_files, 1):
fpath = os.path.join(raw_dir, fname)
img = cv2.imread(fpath)
if img is None:
continue
boxes = detect_boxes(img, templates)
base_name = os.path.splitext(fname)[0]
# 1. Save YOLO label .txt
txt_path = os.path.join(labels_dir, f"{base_name}.txt")
write_yolo_labels(boxes, img.shape, txt_path)
# 2. Save annotated verification image
annotated_img = draw_annotations(img, boxes)
preview_h = 1080
preview_w = int(img.shape[1] * (preview_h / img.shape[0]))
preview_img = cv2.resize(annotated_img, (preview_w, preview_h), interpolation=cv2.INTER_AREA)
jpg_path = os.path.join(annotated_dir, f"{base_name}.jpg")
cv2.imwrite(jpg_path, preview_img, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
total_boxes += len(boxes)
if len(boxes) > 0:
annotated_images_count += 1
found_classes = ", ".join(sorted(set(CLASSES[b[0]] for b in boxes))) if boxes else "None"
print(f"[{i:02d}/{len(image_files)}] {fname} -> {len(boxes)} boxes ({found_classes})")
# 3. Generate data.yaml
yaml_path = os.path.join(root_dir, "dataset", "data.yaml")
yaml_content = f"""# DirtyLeague YOLOv8 Dataset Configuration
path: {os.path.abspath(os.path.join(root_dir, 'dataset'))}
train: images/train
val: images/val
names:
"""
for idx, cname in enumerate(CLASSES):
yaml_content += f" {idx}: {cname}\n"
with open(yaml_path, "w", encoding="utf-8") as f:
f.write(yaml_content)
print("\n" + "=" * 65)
print(f" AUTO-LABELING COMPLETE!")
print(f" Images processed : {len(image_files)}")
print(f" Images with boxes : {annotated_images_count} / {len(image_files)}")
print(f" Total boxes placed : {total_boxes}")
print(f" Labels directory : {labels_dir}")
print(f" Visual previews : {annotated_dir}")
print(f" Config file : {yaml_path}")
print("=" * 65)
if __name__ == "__main__":
main()

View File

@ -9,11 +9,13 @@ import logging
import os import os
import re import re
import time import time
from typing import Dict, Optional, Tuple from typing import Dict, List, Optional, Tuple
import cv2 import cv2
import numpy as np import numpy as np
from window_utils import WindowManager from window_utils import WindowManager
from ocr_utils import ocr_reader from ocr_utils import ocr_reader
from yolo_detector import YoloDetector, Detection
from card_embedder import CardEmbedder
class GameState(Enum): class GameState(Enum):
@ -47,6 +49,23 @@ class DirtyLeagueBot:
REL_ROI_OFFER_CLOSE = (0.8854, 0.0, 1.0, 0.1463) REL_ROI_OFFER_CLOSE = (0.8854, 0.0, 1.0, 0.1463)
REL_CROWN_CHEST_POS = (0.2279, 0.6146) REL_CROWN_CHEST_POS = (0.2279, 0.6146)
# 5 Player card slots (left vertical column) & 5 Opponent card slots (right vertical column)
PLAYER_CARD_SLOTS = [
(0.226, 0.150, 0.294, 0.278),
(0.226, 0.295, 0.294, 0.423),
(0.226, 0.440, 0.294, 0.568),
(0.226, 0.585, 0.294, 0.713),
(0.226, 0.730, 0.294, 0.858),
]
OPPONENT_CARD_SLOTS = [
(0.693, 0.150, 0.761, 0.278),
(0.693, 0.295, 0.761, 0.423),
(0.693, 0.440, 0.761, 0.568),
(0.693, 0.585, 0.761, 0.713),
(0.693, 0.730, 0.761, 0.858),
]
def __init__(self, config: dict): def __init__(self, config: dict):
self.config = config self.config = config
self.bot_conf = config.get("bot", {}) self.bot_conf = config.get("bot", {})
@ -78,6 +97,10 @@ class DirtyLeagueBot:
process_name=self.target_conf.get("process_name", "DirtyLeague.exe") 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 # Template storage with dynamic scaling cache
self.base_templates: Dict[str, np.ndarray] = {} self.base_templates: Dict[str, np.ndarray] = {}
self.scaled_templates: Dict[str, np.ndarray] = {} self.scaled_templates: Dict[str, np.ndarray] = {}
@ -254,9 +277,44 @@ class DirtyLeagueBot:
return True return True
return False return False
def detect_state(self, frame: np.ndarray) -> GameState: def detect_state(
"""Determines current game screen based on visible templates (accelerated with relative ROIs).""" self, frame: np.ndarray, dets: Optional[Dict[str, List[Detection]]] = None
# 1. Ultra-fast relative ROI checks (~3-5ms each on any resolution) ) -> 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 \ 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): self.find_template_in_rel_roi(frame, "btn_open_chest", self.REL_ROI_NO_SLOTS_OPEN):
return GameState.NO_FREE_SLOTS_SCREEN return GameState.NO_FREE_SLOTS_SCREEN
@ -280,7 +338,7 @@ class DirtyLeagueBot:
self.find_template_in_rel_roi(frame, "banner_victory", self.REL_ROI_VICTORY_BANNER, threshold=0.65): self.find_template_in_rel_roi(frame, "banner_victory", self.REL_ROI_VICTORY_BANNER, threshold=0.65):
return GameState.VICTORY_SCREEN 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"): if self.find_template(frame, "btn_remove_chest") or self.find_template(frame, "btn_open_chest"):
return GameState.NO_FREE_SLOTS_SCREEN 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"): if self.find_template(frame, "text_victory") or self.find_template(frame, "banner_victory") or self.find_template(frame, "btn_collect_victory"):
@ -359,11 +417,20 @@ class DirtyLeagueBot:
return results 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). 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"]: 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) pos = self.find_template_in_rel_roi(frame, tpl_name, self.REL_ROI_OFFER_CLOSE, threshold=0.55)
if pos: if pos:
@ -373,9 +440,18 @@ class DirtyLeagueBot:
return True return True
return False 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.") logging.info("State: CHEST_OPEN_SCREEN.")
# Step 1: Check if "Turn all cards over" is visible # 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 \ 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") self.find_template(frame, "btn_turn_all_cards")
if turn_pos: if turn_pos:
@ -385,6 +461,15 @@ class DirtyLeagueBot:
return return
# Step 2: Check if "Collect" is visible # 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 \ collect_pos = self.find_template_in_rel_roi(frame, "btn_collect", self.REL_ROI_CHEST_COLLECT) or \
self.find_template(frame, "btn_collect") self.find_template(frame, "btn_collect")
if collect_pos: if collect_pos:
@ -422,8 +507,7 @@ class DirtyLeagueBot:
def fast_surrender_pipeline(self, max_wait_sec: float = 15.0) -> bool: def fast_surrender_pipeline(self, max_wait_sec: float = 15.0) -> bool:
""" """
Ultra-fast surrender pipeline for FORCED DERANK. Ultra-fast surrender pipeline for FORCED DERANK.
Polls Exit button relative ROI (~20ms per check). The millisecond battle finishes loading, Uses YOLO detection to detect Exit / Leave buttons instantly, falling back to relative ROI.
clicks Exit via fast click, then immediately polls Leave dialog relative ROI and confirms.
Surrenders the battle within < 0.15s of battle load, before auto-battle can deal lethal damage. 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...") logging.info("[FAST DERANK] High-frequency surrender monitor engaged. Waiting for battle screen...")
@ -439,6 +523,20 @@ class DirtyLeagueBot:
time.sleep(0.02) time.sleep(0.02)
continue 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) exit_pos = self.find_template_in_rel_roi(frame, "btn_exit", self.REL_ROI_EXIT, threshold=0.75)
if exit_pos: if exit_pos:
logging.info(f"[FAST DERANK] Battle loaded! 'Exit' detected at {exit_pos} in {time.time() - t0:.2f}s. Clicking...") logging.info(f"[FAST DERANK] Battle loaded! 'Exit' detected at {exit_pos} in {time.time() - t0:.2f}s. Clicking...")
@ -468,6 +566,14 @@ class DirtyLeagueBot:
time.sleep(0.02) time.sleep(0.02)
continue 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) leave_pos = self.find_template_in_rel_roi(frame, "btn_leave", self.REL_ROI_LEAVE, threshold=0.75)
if leave_pos: if leave_pos:
self.click_client_pos_fast(leave_pos[0], leave_pos[1]) self.click_client_pos_fast(leave_pos[0], leave_pos[1])
@ -480,7 +586,7 @@ class DirtyLeagueBot:
logging.warning("[FAST DERANK] Leave button not found after clicking Exit.") logging.warning("[FAST DERANK] Leave button not found after clicking Exit.")
return False 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.") logging.info("State: TOWER_LOBBY.")
self._result_recorded_for_match = False self._result_recorded_for_match = False
fh, fw = frame.shape[:2] fh, fw = frame.shape[:2]
@ -502,8 +608,11 @@ class DirtyLeagueBot:
logging.info(f"[CROWN CHEST] Progress: {c_current}/{c_max} [CROWN]") logging.info(f"[CROWN CHEST] Progress: {c_current}/{c_max} [CROWN]")
if c_current >= c_max: if c_current >= c_max:
logging.info(f"[CROWN CHEST] Goal reached ({c_current}/{c_max})! Clicking crown chest...") logging.info(f"[CROWN CHEST] Goal reached ({c_current}/{c_max})! Clicking crown chest...")
cx = int(self.REL_CROWN_CHEST_POS[0] * fw) if dets and "crown_chest" in dets:
cy = int(self.REL_CROWN_CHEST_POS[1] * fh) 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) self.click_client_pos(cx, cy)
time.sleep(1.5) time.sleep(1.5)
return return
@ -521,9 +630,14 @@ class DirtyLeagueBot:
return return
logging.info("Searching for 'btn_fight' / 'btn_play'...") logging.info("Searching for 'btn_fight' / 'btn_play'...")
fight_pos = self.find_template_in_rel_roi(frame, "btn_fight", self.REL_ROI_FIGHT) or \ if dets and "btn_fight" in dets:
self.find_template(frame, "btn_fight") or \ fight_pos = dets["btn_fight"][0].center
self.find_template(frame, "btn_play") 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: if fight_pos:
self.click_client_pos(fight_pos[0], fight_pos[1]) self.click_client_pos(fight_pos[0], fight_pos[1])
logging.info("Clicked fight button. Transitioning towards IN_GAME.") logging.info("Clicked fight button. Transitioning towards IN_GAME.")
@ -532,7 +646,24 @@ class DirtyLeagueBot:
else: else:
time.sleep(1.5) time.sleep(1.5)
def handle_in_game(self, frame: np.ndarray): def extract_battle_cards(self, frame: np.ndarray) -> Tuple[List[np.ndarray], List[np.ndarray]]:
"""Extracts exact 5 player and 5 opponent card portrait crops from battle screen."""
fh, fw = frame.shape[:2]
player_crops = []
for rx1, ry1, rx2, ry2 in self.PLAYER_CARD_SLOTS:
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
player_crops.append(frame[y1:y2, x1:x2])
opponent_crops = []
for rx1, ry1, rx2, ry2 in self.OPPONENT_CARD_SLOTS:
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
opponent_crops.append(frame[y1:y2, x1:x2])
return player_crops, opponent_crops
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) 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})" reason = "Forced Derank" if self.derank_mode else f"Win Streak Limit ({self.current_wins}/{self.max_wins_limit})"
@ -541,40 +672,70 @@ class DirtyLeagueBot:
self.fast_surrender_pipeline(max_wait_sec=5.0) self.fast_surrender_pipeline(max_wait_sec=5.0)
return return
else: else:
player_crops, opponent_crops = self.extract_battle_cards(frame)
player_names = []
for i, crop in enumerate(player_crops, 1):
name, sim = self.embedder.identify_card(crop)
player_names.append(f"#{i}:{name or 'Unknown'}({sim:.2f})")
enemy_names = []
for i, crop in enumerate(opponent_crops, 1):
name, sim = self.embedder.identify_card(crop)
enemy_names.append(f"#{i}:{name or 'Unknown'}({sim:.2f})")
logging.info(f"[IN_GAME] Player Team: [{', '.join(player_names)}] vs Opponent: [{', '.join(enemy_names)}]")
logging.info("Match in progress (Normal Play, Auto-battle active). Waiting for match results...") 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: if not self._result_recorded_for_match:
self._result_recorded_for_match = True self._result_recorded_for_match = True
self.current_wins += 1 self.current_wins += 1
self.total_games += 1 self.total_games += 1
logging.info(f"[WIN] Victory recorded! Streak: {self.current_wins}/{self.max_wins_limit}. Total games: {self.total_games}") 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 \ if dets and "btn_collect" in dets:
self.find_template(frame, "btn_collect_victory") pos = dets["btn_collect"][0].center
if collect_pos: logging.info(f"[WIN] Clicking 'Collect' via YOLO at {pos}")
self.click_client_pos(collect_pos[0], collect_pos[1]) self.click_client_pos(pos[0], pos[1])
else: 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) 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: if not self._result_recorded_for_match:
self._result_recorded_for_match = True self._result_recorded_for_match = True
self.current_wins = 0 self.current_wins = 0
self.total_games += 1 self.total_games += 1
logging.info(f"[LOSS] Streak reset to 0. Total games: {self.total_games}") 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 \ if dets and "btn_ok" in dets:
self.find_template(frame, "btn_ok") pos = dets["btn_ok"][0].center
if ok_pos: logging.info(f"[DEFEAT] Clicking 'OK' via YOLO at {pos}")
self.click_client_pos(ok_pos[0], ok_pos[1]) self.click_client_pos(pos[0], pos[1])
else: 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) 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').") 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 \ 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") self.find_template(frame, "btn_open_chest")
if open_pos: if open_pos:
@ -590,12 +751,13 @@ class DirtyLeagueBot:
def step(self): def step(self):
"""Executes a single FSM iteration.""" """Executes a single FSM iteration."""
frame = self.wm.capture_frame(client_only=True) 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 # 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 return
detected_state = self.detect_state(frame) detected_state = self.detect_state(frame, dets)
if detected_state != self.state: if detected_state != self.state:
logging.info(f"State changed: {self.state.value} -> {detected_state.value}") logging.info(f"State changed: {self.state.value} -> {detected_state.value}")
@ -604,17 +766,17 @@ class DirtyLeagueBot:
if self.state == GameState.MAIN_MENU: if self.state == GameState.MAIN_MENU:
self.handle_main_menu(frame) self.handle_main_menu(frame)
elif self.state == GameState.TOWER_LOBBY: elif self.state == GameState.TOWER_LOBBY:
self.handle_lobby(frame) self.handle_lobby(frame, dets)
elif self.state == GameState.CHEST_OPEN_SCREEN: 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: elif self.state == GameState.IN_GAME:
self.handle_in_game(frame) self.handle_in_game(frame, dets)
elif self.state == GameState.VICTORY_SCREEN: elif self.state == GameState.VICTORY_SCREEN:
self.handle_victory(frame) self.handle_victory(frame, dets)
elif self.state == GameState.DEFEAT_SCREEN: elif self.state == GameState.DEFEAT_SCREEN:
self.handle_defeat(frame) self.handle_defeat(frame, dets)
elif self.state == GameState.NO_FREE_SLOTS_SCREEN: 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: if self.state == GameState.UNKNOWN:
self._unknown_state_count += 1 self._unknown_state_count += 1
if self._unknown_state_count % 15 == 0: 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

24
collect_data.bat Normal file
View File

@ -0,0 +1,24 @@
@echo off
chcp 65001 > nul
echo ========================================================
echo DirtyLeague Dataset Collector for YOLO
echo ========================================================
echo.
echo Modes:
echo 1. Auto (captures on screen transitions) [DEFAULT]
echo 2. Manual (press 'C' to save current screen)
echo.
echo Press 'Q' at any time to finish collection.
echo.
if exist ".venv\Scripts\python.exe" (
".venv\Scripts\python.exe" collect_dataset.py %*
) else (
python collect_dataset.py %*
)
if %ERRORLEVEL% neq 0 (
echo.
echo Collector stopped with exit code %ERRORLEVEL%.
pause
)

141
collect_dataset.py Normal file
View File

@ -0,0 +1,141 @@
"""
Dataset Collection Tool for DirtyLeague.
Captures diverse game screens for training YOLO object detection and Card Embeddings.
Modes:
- 'auto': Automatically captures new frames when significant screen changes occur (diff > threshold).
- 'manual': Captures frame upon pressing 'C'.
- 'interval': Captures frame every N seconds.
Press 'Q' at any time to finish collection.
"""
import argparse
import os
import sys
import time
import cv2
import numpy as np
import keyboard
from window_utils import WindowManager, ensure_interactive_desktop
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "dataset", "raw")
def frame_diff_ratio(img1: np.ndarray, img2: np.ndarray) -> float:
"""Computes mean normalized absolute difference between two frames."""
if img1.shape != img2.shape:
return 1.0
# Downscale for fast difference calculation
small1 = cv2.resize(img1, (320, 180), interpolation=cv2.INTER_AREA)
small2 = cv2.resize(img2, (320, 180), interpolation=cv2.INTER_AREA)
diff = cv2.absdiff(small1, small2)
return float(np.mean(diff)) / 255.0
def main():
parser = argparse.ArgumentParser(description="Collect game screens for YOLO / Embedding dataset")
parser.add_argument("--mode", choices=["auto", "manual", "interval"], default="auto",
help="Capture mode: 'auto' (on screen changes), 'manual' (press C), 'interval' (every N sec)")
parser.add_argument("--interval", type=float, default=1.5, help="Sampling interval in seconds (default: 1.5)")
parser.add_argument("--diff-threshold", type=float, default=0.04,
help="Minimum change ratio to trigger auto save (default: 0.04 = 4%%)")
parser.add_argument("--max-count", type=int, default=100, help="Target number of screenshots to collect (default: 100)")
parser.add_argument("--out-dir", type=str, default=OUTPUT_DIR, help="Directory to save raw screenshots")
args = parser.parse_args()
os.makedirs(args.out_dir, exist_ok=True)
existing_count = len([f for f in os.listdir(args.out_dir) if f.endswith(".png")])
print("=" * 65)
print(" DIRTYLEAGUE DATASET COLLECTION TOOL")
print("=" * 65)
print(f"Output directory : {args.out_dir}")
print(f"Existing images : {existing_count}")
print(f"Target count : {args.max_count}")
print(f"Capture mode : {args.mode.upper()}")
if args.mode == "auto":
print(f"Auto-trigger : Screen diff > {args.diff_threshold * 100:.1f}%, check every {args.interval}s")
elif args.mode == "manual":
print("Manual hotkey : Press 'C' to capture screen")
elif args.mode == "interval":
print(f"Interval trigger : Every {args.interval}s")
print("Exit hotkey : Press 'Q' or Ctrl+C to stop")
print("=" * 65)
wm = WindowManager()
hwnd = wm.find_window()
if not hwnd:
print("[ERROR] DirtyLeague window not found! Please ensure the game is running.")
sys.exit(1)
saved_count = 0
last_saved_frame = None
last_capture_time = 0.0
print("\n[READY] Collection started. Switch to the game and play/navigate screens...")
try:
while saved_count < args.max_count:
if keyboard.is_pressed("q"):
print("\n[INFO] 'Q' pressed. Stopping collection...")
break
now = time.time()
frame = wm.capture_frame(client_only=True)
if frame is None or frame.size == 0:
time.sleep(0.5)
continue
should_save = False
reason = ""
if args.mode == "manual":
if keyboard.is_pressed("c"):
should_save = True
reason = "Manual (C key)"
time.sleep(0.3) # Debounce keypress
elif args.mode == "interval":
if now - last_capture_time >= args.interval:
should_save = True
reason = f"Interval ({args.interval}s)"
elif args.mode == "auto":
if last_saved_frame is None:
should_save = True
reason = "Initial frame"
elif now - last_capture_time >= args.interval:
diff = frame_diff_ratio(frame, last_saved_frame)
if diff >= args.diff_threshold:
should_save = True
reason = f"Screen changed ({diff * 100:.1f}% >= {args.diff_threshold * 100:.1f}%)"
if should_save:
ts = time.strftime("%Y%m%d_%H%M%S")
millis = int((now % 1) * 1000)
filename = f"dl_screen_{ts}_{millis:03d}.png"
filepath = os.path.join(args.out_dir, filename)
cv2.imwrite(filepath, frame)
saved_count += 1
last_saved_frame = frame.copy()
last_capture_time = now
h, w = frame.shape[:2]
print(f"[{saved_count:03d}/{args.max_count}] Saved: {filename} ({w}x{h}) | Reason: {reason}")
time.sleep(0.1)
except KeyboardInterrupt:
print("\n[INFO] Interrupted by user.")
total_in_dir = len([f for f in os.listdir(args.out_dir) if f.endswith(".png")])
print("\n" + "=" * 65)
print(f" COLLECTION FINISHED: {saved_count} new image(s) saved.")
print(f" Total images in '{args.out_dir}': {total_in_dir}")
print("=" * 65)
if __name__ == "__main__":
main()

17
dataset/data.yaml Normal file
View File

@ -0,0 +1,17 @@
# DirtyLeague YOLOv8 Dataset Configuration
path: F:/Gitea/DirtyLeague_Bot/dataset
train: images/train
val: images/val
names:
0: btn_fight
1: btn_exit
2: btn_leave
3: btn_ok
4: btn_collect
5: btn_turn_all
6: btn_close
7: btn_remove_chest
8: btn_open_chest
9: crown_chest
10: card

View File

@ -0,0 +1,2 @@
0 0.852469 0.425676 0.116821 0.080370
9 0.227619 0.614509 0.119631 0.139403

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.955610 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,6 @@
1 0.459635 0.956098 0.052083 0.024390
2 0.554427 0.628780 0.079167 0.030244
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,2 @@
0 0.853516 0.425610 0.110677 0.080488
9 0.227865 0.614146 0.119792 0.139512

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,5 @@
1 0.459635 0.956098 0.052083 0.024390
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1 @@
4 0.500000 0.715854 0.148438 0.056098

View File

@ -0,0 +1 @@
5 0.500000 0.715854 0.148438 0.056098

View File

@ -0,0 +1 @@
5 0.500000 0.715854 0.148438 0.056098

View File

@ -0,0 +1 @@
3 0.304297 0.775854 0.095052 0.044390

View File

@ -0,0 +1,6 @@
1 0.459635 0.956098 0.052083 0.024390
2 0.554427 0.628780 0.079167 0.030244
10 0.219792 0.849756 0.036458 0.092683
10 0.319792 0.849756 0.036458 0.092683
10 0.679948 0.849756 0.036458 0.092683
10 0.779948 0.849756 0.036458 0.092683

View File

@ -0,0 +1,2 @@
7 0.431901 0.865610 0.117448 0.044390
8 0.568099 0.865610 0.117448 0.044390

View File

@ -0,0 +1,2 @@
6 0.965885 0.060976 0.017708 0.033171
6 0.965885 0.060976 0.017708 0.033171

View File

@ -0,0 +1 @@
4 0.250260 0.804634 0.095312 0.044390

1
dataset/raw/.gitkeep Normal file
View File

@ -0,0 +1 @@
# raw dataset directory

110
inspect_battle_cards.py Normal file
View File

@ -0,0 +1,110 @@
"""
Inspection script to calibrate exact bounding boxes for:
- 5 Player cards (left vertical column)
- 5 Opponent cards (right vertical column)
on the DirtyLeague Battle Screen (3840x2050 design resolution).
"""
import os
import cv2
import numpy as np
from card_embedder import CardEmbedder
# Normalized relative coordinates (rx1, ry1, rx2, ry2) for 3840x2050
# Calibrated from actual 4K battle screenshot
PLAYER_CARD_SLOTS = [
# Slot 1 (top left)
(0.226, 0.150, 0.294, 0.278),
# Slot 2
(0.226, 0.295, 0.294, 0.423),
# Slot 3
(0.226, 0.440, 0.294, 0.568),
# Slot 4
(0.226, 0.585, 0.294, 0.713),
# Slot 5 (bottom left)
(0.226, 0.730, 0.294, 0.858),
]
OPPONENT_CARD_SLOTS = [
# Slot 1 (top right)
(0.693, 0.150, 0.761, 0.278),
# Slot 2
(0.693, 0.295, 0.761, 0.423),
# Slot 3
(0.693, 0.440, 0.761, 0.568),
# Slot 4
(0.693, 0.585, 0.761, 0.713),
# Slot 5 (bottom right)
(0.693, 0.730, 0.761, 0.858),
]
def extract_cards(frame: np.ndarray):
fh, fw = frame.shape[:2]
player_crops = []
opponent_crops = []
for rx1, ry1, rx2, ry2 in PLAYER_CARD_SLOTS:
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
player_crops.append(frame[y1:y2, x1:x2])
for rx1, ry1, rx2, ry2 in OPPONENT_CARD_SLOTS:
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
opponent_crops.append(frame[y1:y2, x1:x2])
return player_crops, opponent_crops
def main():
img_path = os.path.join("dataset", "raw", "dl_screen_20260910_233743_310.png")
if not os.path.exists(img_path):
print(f"[ERROR] Sample battle screenshot not found at: {img_path}")
return
frame = cv2.imread(img_path)
fh, fw = frame.shape[:2]
print(f"\nAnalyzing battle frame resolution: {fw}x{fh}")
player_crops, opponent_crops = extract_cards(frame)
os.makedirs("debug_cards", exist_ok=True)
embedder = CardEmbedder()
print("\n--- Player Cards (Left Column 1..5) ---")
for i, crop in enumerate(player_crops, 1):
crop_path = os.path.join("debug_cards", f"player_card_{i}.png")
cv2.imwrite(crop_path, crop)
vec = embedder.embed(crop)
print(f" Player Slot #{i}: {crop.shape[1]}x{crop.shape[0]}px | Vector norm: {np.linalg.norm(vec):.4f} | Saved to: {crop_path}")
print("\n--- Opponent Cards (Right Column 1..5) ---")
for i, crop in enumerate(opponent_crops, 1):
crop_path = os.path.join("debug_cards", f"opponent_card_{i}.png")
cv2.imwrite(crop_path, crop)
vec = embedder.embed(crop)
print(f" Opponent Slot #{i}: {crop.shape[1]}x{crop.shape[0]}px | Vector norm: {np.linalg.norm(vec):.4f} | Saved to: {crop_path}")
# Draw visual overlay on whole frame
annotated = frame.copy()
for i, (rx1, ry1, rx2, ry2) in enumerate(PLAYER_CARD_SLOTS, 1):
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 3)
cv2.putText(annotated, f"Player #{i}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)
for i, (rx1, ry1, rx2, ry2) in enumerate(OPPONENT_CARD_SLOTS, 1):
x1, y1 = int(rx1 * fw), int(ry1 * fh)
x2, y2 = int(rx2 * fw), int(ry2 * fh)
cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 0, 255), 3)
cv2.putText(annotated, f"Enemy #{i}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
preview_path = os.path.join("debug_cards", "annotated_battle_slots.png")
cv2.imwrite(preview_path, cv2.resize(annotated, (1920, 1025)))
print(f"\n[SUCCESS] Full annotated battle preview saved to: {preview_path}")
if __name__ == "__main__":
main()

BIN
models/card_embedder.onnx Normal file

Binary file not shown.

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