DirtyLeague_Bot/collect_dataset.py

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