Compare commits
4 Commits
main
...
feature/ve
| Author | SHA1 | Date | |
|---|---|---|---|
| aa4d866d7d | |||
| 3a40c2003f | |||
| 35b0e53413 | |||
| bfeb99b47d |
10
.gitignore
vendored
10
.gitignore
vendored
@ -19,3 +19,13 @@ logs/
|
||||
*.log
|
||||
temp_*.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
259
auto_label.py
Normal 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()
|
||||
238
bot_core.py
238
bot_core.py
@ -9,11 +9,13 @@ import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Dict, Optional, Tuple
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import cv2
|
||||
import numpy as np
|
||||
from window_utils import WindowManager
|
||||
from ocr_utils import ocr_reader
|
||||
from yolo_detector import YoloDetector, Detection
|
||||
from card_embedder import CardEmbedder
|
||||
|
||||
|
||||
class GameState(Enum):
|
||||
@ -47,6 +49,23 @@ class DirtyLeagueBot:
|
||||
REL_ROI_OFFER_CLOSE = (0.8854, 0.0, 1.0, 0.1463)
|
||||
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):
|
||||
self.config = config
|
||||
self.bot_conf = config.get("bot", {})
|
||||
@ -78,6 +97,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 +277,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 +338,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 +417,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 +440,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 +461,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 +507,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 +523,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 +566,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 +586,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 +608,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 +630,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 +646,24 @@ class DirtyLeagueBot:
|
||||
else:
|
||||
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)
|
||||
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)
|
||||
return
|
||||
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...")
|
||||
|
||||
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 +751,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 +766,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
152
card_embedder.py
Normal 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
24
collect_data.bat
Normal 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
141
collect_dataset.py
Normal 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
17
dataset/data.yaml
Normal 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
|
||||
2
dataset/labels/dl_screen_20260909_092820_419.txt
Normal file
2
dataset/labels/dl_screen_20260909_092820_419.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.852469 0.425676 0.116821 0.080370
|
||||
9 0.227619 0.614509 0.119631 0.139403
|
||||
0
dataset/labels/dl_screen_20260910_233619_719.txt
Normal file
0
dataset/labels/dl_screen_20260910_233619_719.txt
Normal file
0
dataset/labels/dl_screen_20260910_233627_942.txt
Normal file
0
dataset/labels/dl_screen_20260910_233627_942.txt
Normal file
0
dataset/labels/dl_screen_20260910_233628_653.txt
Normal file
0
dataset/labels/dl_screen_20260910_233628_653.txt
Normal file
0
dataset/labels/dl_screen_20260910_233634_909.txt
Normal file
0
dataset/labels/dl_screen_20260910_233634_909.txt
Normal file
0
dataset/labels/dl_screen_20260910_233635_503.txt
Normal file
0
dataset/labels/dl_screen_20260910_233635_503.txt
Normal file
0
dataset/labels/dl_screen_20260910_233639_307.txt
Normal file
0
dataset/labels/dl_screen_20260910_233639_307.txt
Normal file
0
dataset/labels/dl_screen_20260910_233641_883.txt
Normal file
0
dataset/labels/dl_screen_20260910_233641_883.txt
Normal file
0
dataset/labels/dl_screen_20260910_233643_432.txt
Normal file
0
dataset/labels/dl_screen_20260910_233643_432.txt
Normal file
0
dataset/labels/dl_screen_20260910_233645_984.txt
Normal file
0
dataset/labels/dl_screen_20260910_233645_984.txt
Normal file
0
dataset/labels/dl_screen_20260910_233647_086.txt
Normal file
0
dataset/labels/dl_screen_20260910_233647_086.txt
Normal file
0
dataset/labels/dl_screen_20260910_233648_664.txt
Normal file
0
dataset/labels/dl_screen_20260910_233648_664.txt
Normal file
0
dataset/labels/dl_screen_20260910_233656_985.txt
Normal file
0
dataset/labels/dl_screen_20260910_233656_985.txt
Normal file
0
dataset/labels/dl_screen_20260910_233701_948.txt
Normal file
0
dataset/labels/dl_screen_20260910_233701_948.txt
Normal file
2
dataset/labels/dl_screen_20260910_233702_650.txt
Normal file
2
dataset/labels/dl_screen_20260910_233702_650.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233708_424.txt
Normal file
2
dataset/labels/dl_screen_20260910_233708_424.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233710_925.txt
Normal file
0
dataset/labels/dl_screen_20260910_233710_925.txt
Normal file
2
dataset/labels/dl_screen_20260910_233719_952.txt
Normal file
2
dataset/labels/dl_screen_20260910_233719_952.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233723_738.txt
Normal file
2
dataset/labels/dl_screen_20260910_233723_738.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233725_447.txt
Normal file
0
dataset/labels/dl_screen_20260910_233725_447.txt
Normal file
2
dataset/labels/dl_screen_20260910_233732_097.txt
Normal file
2
dataset/labels/dl_screen_20260910_233732_097.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233738_553.txt
Normal file
2
dataset/labels/dl_screen_20260910_233738_553.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233740_071.txt
Normal file
0
dataset/labels/dl_screen_20260910_233740_071.txt
Normal file
5
dataset/labels/dl_screen_20260910_233741_601.txt
Normal file
5
dataset/labels/dl_screen_20260910_233741_601.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233743_310.txt
Normal file
5
dataset/labels/dl_screen_20260910_233743_310.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233745_861.txt
Normal file
5
dataset/labels/dl_screen_20260910_233745_861.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233746_421.txt
Normal file
5
dataset/labels/dl_screen_20260910_233746_421.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233748_121.txt
Normal file
5
dataset/labels/dl_screen_20260910_233748_121.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233750_285.txt
Normal file
5
dataset/labels/dl_screen_20260910_233750_285.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233751_818.txt
Normal file
5
dataset/labels/dl_screen_20260910_233751_818.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233753_339.txt
Normal file
5
dataset/labels/dl_screen_20260910_233753_339.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233754_842.txt
Normal file
5
dataset/labels/dl_screen_20260910_233754_842.txt
Normal 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
|
||||
0
dataset/labels/dl_screen_20260910_233756_437.txt
Normal file
0
dataset/labels/dl_screen_20260910_233756_437.txt
Normal file
0
dataset/labels/dl_screen_20260910_233800_157.txt
Normal file
0
dataset/labels/dl_screen_20260910_233800_157.txt
Normal file
2
dataset/labels/dl_screen_20260910_233801_779.txt
Normal file
2
dataset/labels/dl_screen_20260910_233801_779.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233805_494.txt
Normal file
2
dataset/labels/dl_screen_20260910_233805_494.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233807_025.txt
Normal file
0
dataset/labels/dl_screen_20260910_233807_025.txt
Normal file
2
dataset/labels/dl_screen_20260910_233813_319.txt
Normal file
2
dataset/labels/dl_screen_20260910_233813_319.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233822_125.txt
Normal file
2
dataset/labels/dl_screen_20260910_233822_125.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233823_640.txt
Normal file
0
dataset/labels/dl_screen_20260910_233823_640.txt
Normal file
5
dataset/labels/dl_screen_20260910_233825_207.txt
Normal file
5
dataset/labels/dl_screen_20260910_233825_207.txt
Normal 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
|
||||
6
dataset/labels/dl_screen_20260910_233827_908.txt
Normal file
6
dataset/labels/dl_screen_20260910_233827_908.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233828_534.txt
Normal file
5
dataset/labels/dl_screen_20260910_233828_534.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233830_047.txt
Normal file
5
dataset/labels/dl_screen_20260910_233830_047.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233831_574.txt
Normal file
5
dataset/labels/dl_screen_20260910_233831_574.txt
Normal 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
|
||||
0
dataset/labels/dl_screen_20260910_233833_274.txt
Normal file
0
dataset/labels/dl_screen_20260910_233833_274.txt
Normal file
0
dataset/labels/dl_screen_20260910_233835_865.txt
Normal file
0
dataset/labels/dl_screen_20260910_233835_865.txt
Normal file
0
dataset/labels/dl_screen_20260910_233837_764.txt
Normal file
0
dataset/labels/dl_screen_20260910_233837_764.txt
Normal file
2
dataset/labels/dl_screen_20260910_233839_343.txt
Normal file
2
dataset/labels/dl_screen_20260910_233839_343.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
2
dataset/labels/dl_screen_20260910_233842_857.txt
Normal file
2
dataset/labels/dl_screen_20260910_233842_857.txt
Normal file
@ -0,0 +1,2 @@
|
||||
0 0.853516 0.425610 0.110677 0.080488
|
||||
9 0.227865 0.614146 0.119792 0.139512
|
||||
0
dataset/labels/dl_screen_20260910_233843_378.txt
Normal file
0
dataset/labels/dl_screen_20260910_233843_378.txt
Normal file
5
dataset/labels/dl_screen_20260910_233845_895.txt
Normal file
5
dataset/labels/dl_screen_20260910_233845_895.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233846_656.txt
Normal file
5
dataset/labels/dl_screen_20260910_233846_656.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233848_158.txt
Normal file
5
dataset/labels/dl_screen_20260910_233848_158.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233849_674.txt
Normal file
5
dataset/labels/dl_screen_20260910_233849_674.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233854_854.txt
Normal file
5
dataset/labels/dl_screen_20260910_233854_854.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233856_369.txt
Normal file
5
dataset/labels/dl_screen_20260910_233856_369.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233858_891.txt
Normal file
5
dataset/labels/dl_screen_20260910_233858_891.txt
Normal 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
|
||||
5
dataset/labels/dl_screen_20260910_233859_411.txt
Normal file
5
dataset/labels/dl_screen_20260910_233859_411.txt
Normal 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
|
||||
0
dataset/labels/dl_screen_20260910_233901_081.txt
Normal file
0
dataset/labels/dl_screen_20260910_233901_081.txt
Normal file
0
dataset/labels/dl_screen_20260910_233904_508.txt
Normal file
0
dataset/labels/dl_screen_20260910_233904_508.txt
Normal file
1
dataset/labels/dl_screen_chest_collect_4k.txt
Normal file
1
dataset/labels/dl_screen_chest_collect_4k.txt
Normal file
@ -0,0 +1 @@
|
||||
4 0.500000 0.715854 0.148438 0.056098
|
||||
1
dataset/labels/dl_screen_chest_reward_4k.txt
Normal file
1
dataset/labels/dl_screen_chest_reward_4k.txt
Normal file
@ -0,0 +1 @@
|
||||
5 0.500000 0.715854 0.148438 0.056098
|
||||
1
dataset/labels/dl_screen_crown_chest_open_4k.txt
Normal file
1
dataset/labels/dl_screen_crown_chest_open_4k.txt
Normal file
@ -0,0 +1 @@
|
||||
5 0.500000 0.715854 0.148438 0.056098
|
||||
1
dataset/labels/dl_screen_defeat_screen_4k.txt
Normal file
1
dataset/labels/dl_screen_defeat_screen_4k.txt
Normal file
@ -0,0 +1 @@
|
||||
3 0.304297 0.775854 0.095052 0.044390
|
||||
6
dataset/labels/dl_screen_leave_dialog_4k.txt
Normal file
6
dataset/labels/dl_screen_leave_dialog_4k.txt
Normal 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
|
||||
2
dataset/labels/dl_screen_no_free_slots_4k.txt
Normal file
2
dataset/labels/dl_screen_no_free_slots_4k.txt
Normal file
@ -0,0 +1,2 @@
|
||||
7 0.431901 0.865610 0.117448 0.044390
|
||||
8 0.568099 0.865610 0.117448 0.044390
|
||||
2
dataset/labels/dl_screen_offer_popup_4k.txt
Normal file
2
dataset/labels/dl_screen_offer_popup_4k.txt
Normal file
@ -0,0 +1,2 @@
|
||||
6 0.965885 0.060976 0.017708 0.033171
|
||||
6 0.965885 0.060976 0.017708 0.033171
|
||||
1
dataset/labels/dl_screen_victory_screen_4k.txt
Normal file
1
dataset/labels/dl_screen_victory_screen_4k.txt
Normal file
@ -0,0 +1 @@
|
||||
4 0.250260 0.804634 0.095312 0.044390
|
||||
1
dataset/raw/.gitkeep
Normal file
1
dataset/raw/.gitkeep
Normal file
@ -0,0 +1 @@
|
||||
# raw dataset directory
|
||||
110
inspect_battle_cards.py
Normal file
110
inspect_battle_cards.py
Normal 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
BIN
models/card_embedder.onnx
Normal file
Binary file not shown.
BIN
models/dirty_league_yolo.onnx
Normal file
BIN
models/dirty_league_yolo.onnx
Normal file
Binary file not shown.
60
test_vector_vision.py
Normal file
60
test_vector_vision.py
Normal 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
159
train_yolo.py
Normal 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
230
yolo_detector.py
Normal 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
|
||||
Reference in New Issue
Block a user