From bfeb99b47d1c72af031eea6e5092680351aa3c80 Mon Sep 17 00:00:00 2001 From: VolandSZ <{E-MAIL}> Date: Wed, 9 Sep 2026 09:28:57 +0300 Subject: [PATCH] Add dataset collector tool for YOLO object detection and card embeddings --- .gitignore | 1 + collect_data.bat | 24 ++++++++ collect_dataset.py | 141 +++++++++++++++++++++++++++++++++++++++++++ dataset/raw/.gitkeep | 1 + 4 files changed, 167 insertions(+) create mode 100644 collect_data.bat create mode 100644 collect_dataset.py create mode 100644 dataset/raw/.gitkeep diff --git a/.gitignore b/.gitignore index 16df885..c42a48a 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ logs/ *.log temp_*.png debug_*.png +dataset/raw/*.png diff --git a/collect_data.bat b/collect_data.bat new file mode 100644 index 0000000..78a3f73 --- /dev/null +++ b/collect_data.bat @@ -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 +) diff --git a/collect_dataset.py b/collect_dataset.py new file mode 100644 index 0000000..36da7fb --- /dev/null +++ b/collect_dataset.py @@ -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() diff --git a/dataset/raw/.gitkeep b/dataset/raw/.gitkeep new file mode 100644 index 0000000..67e7e49 --- /dev/null +++ b/dataset/raw/.gitkeep @@ -0,0 +1 @@ +# raw dataset directory