{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chess Piece Detection Model Trainer\n",
    "\n",
    "This notebook consolidates the entire workflow for training a YOLOv8-based chess piece detection model. \n",
    "It handles dependency installation, asset downloading, synthetic data generation, training, and ONNX export."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1: Setup\n",
    "# Install system dependencies\n",
    "# !sudo apt-get update && sudo apt-get install -y libcairo2-dev libpango1.0-dev\n",
    "\n",
    "# Install Python packages\n",
    "!pip install ultralytics --upgrade\n",
    "!pip install aiohttp cairosvg gitpython tqdm opencv-python pyyaml\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import subprocess\n",
    "import torch\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "import logging\n",
    "\n",
    "# Configure logging\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format=\"%(asctime)s - %(levelname)s - %(message)s\",\n",
    "    handlers=[logging.StreamHandler(sys.stdout)],\n",
    "    force=True\n",
    ")\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Set up environment\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'\n",
    "WORKING_DIR = os.getcwd()\n",
    "\n",
    "# Check GPU\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f\"Using device: {DEVICE}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    num_gpus = torch.cuda.device_count()\n",
    "    print(f\"Number of GPUs available: {num_gpus}\")\n",
    "    for i in range(num_gpus):\n",
    "        print(f\"  GPU {i}: {torch.cuda.get_device_name(i)}\")\n",
    "else:\n",
    "    print(\"No GPU available\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2: Asset Download\n",
    "# Embeds logic from training/get_pieces.py\n",
    "\n",
    "import asyncio\n",
    "import aiohttp\n",
    "import cairosvg\n",
    "import git\n",
    "import stat\n",
    "import os\n",
    "from pathlib import Path\n",
    "from tqdm.asyncio import tqdm as tqdm_async\n",
    "from tqdm import tqdm\n",
    "from typing import List, Tuple, Optional\n",
    "from concurrent.futures import ProcessPoolExecutor\n",
    "import functools\n",
    "\n",
    "# Constants\n",
    "CHESSCOM_BASE_URL = \"https://images.chesscomfiles.com/chess-themes/pieces\"\n",
    "DEFAULT_OUTPUT_DIR = Path(\"./assets/pieces\")\n",
    "COLORS = [\"w\", \"b\"]\n",
    "PIECES = [\"k\", \"q\", \"r\", \"b\", \"n\", \"p\"]\n",
    "\n",
    "CHESSCOM_PIECE_THEMES = [\n",
    "    \"neo\", \"game_room\", \"wood\", \"glass\", \"gothic\", \"classic\", \"metal\", \"bases\",\n",
    "    \"neo_wood\", \"icy_sea\", \"club\", \"ocean\", \"newspaper\", \"blindfold\", \"space\",\n",
    "    \"cases\", \"condal\", \"3d_chesskid\", \"8_bit\", \"marble\", \"book\", \"alpha\",\n",
    "    \"bubblegum\", \"dash\", \"graffiti\", \"light\", \"lolz\", \"luca\", \"maya\", \"modern\",\n",
    "    \"nature\", \"neon\", \"sky\", \"tigers\", \"tournament\", \"vintage\", \"3d_wood\",\n",
    "    \"3d_staunton\", \"3d_plastic\", \"real_3d\",\n",
    "]\n",
    "\n",
    "LICHESS_REPO_URL = \"https://github.com/lichess-org/lila.git\"\n",
    "\n",
    "async def download_chesscom_piece(\n",
    "    session: aiohttp.ClientSession,\n",
    "    set_name: str,\n",
    "    color: str,\n",
    "    piece: str,\n",
    "    output_dir: Path,\n",
    "    retries: int = 3\n",
    ") -> bool:\n",
    "    url = f\"{CHESSCOM_BASE_URL}/{set_name}/150/{color}{piece}.png\"\n",
    "    save_path = output_dir / set_name / f\"{color}{piece}.png\"\n",
    "    \n",
    "    if save_path.exists():\n",
    "        return True\n",
    "\n",
    "    for attempt in range(retries):\n",
    "        try:\n",
    "            async with session.get(url) as response:\n",
    "                if response.status == 200:\n",
    "                    save_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "                    content = await response.read()\n",
    "                    save_path.write_bytes(content)\n",
    "                    return True\n",
    "                elif response.status == 404:\n",
    "                    return False\n",
    "                else:\n",
    "                    pass # Retry silently\n",
    "        except Exception:\n",
    "            pass\n",
    "\n",
    "        if attempt < retries - 1:\n",
    "            await asyncio.sleep(0.5 * (attempt + 1))\n",
    "\n",
    "    return False\n",
    "\n",
    "async def download_chesscom_pieces(output_dir: Path, concurrency: int = 100):\n",
    "    logger.info(f\"Starting Chess.com downloads to {output_dir}...\")\n",
    "\n",
    "    connector = aiohttp.TCPConnector(limit=concurrency)\n",
    "    timeout = aiohttp.ClientTimeout(total=30)\n",
    "\n",
    "    tasks = []\n",
    "    async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:\n",
    "        for set_name in CHESSCOM_PIECE_THEMES:\n",
    "            for color in COLORS:\n",
    "                for piece in PIECES:\n",
    "                    tasks.append(\n",
    "                        download_chesscom_piece(session, set_name, color, piece, output_dir)\n",
    "                    )\n",
    "\n",
    "        results = []\n",
    "        for coro in tqdm_async.as_completed(tasks, desc=\"Downloading Chess.com pieces\"):\n",
    "            result = await coro\n",
    "            results.append(result)\n",
    "\n",
    "    success_count = sum(results)\n",
    "    logger.info(f\"Chess.com downloads complete: {success_count}/{len(tasks)} successful\")\n",
    "\n",
    "def on_rm_error(func, path, exc_info):\n",
    "    try:\n",
    "        os.chmod(path, stat.S_IWRITE)\n",
    "        func(path)\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Failed to remove {path}: {e}\")\n",
    "\n",
    "def convert_single_svg(args):\n",
    "    \"\"\"Helper function for parallel SVG conversion\"\"\"\n",
    "    svg_path, png_path = args\n",
    "    try:\n",
    "        png_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "        cairosvg.svg2png(\n",
    "            url=str(svg_path),\n",
    "            write_to=str(png_path),\n",
    "            output_height=150,\n",
    "            output_width=150,\n",
    "        )\n",
    "        return True\n",
    "    except Exception:\n",
    "        return False\n",
    "\n",
    "async def download_lichess_pieces(output_dir: Path):\n",
    "    repo_dir = Path(\"temp_lila_repo\")\n",
    "    target_subfolder = \"public/piece\"\n",
    "\n",
    "    logger.info(\"Cloning Lichess repo (this may take a moment)...\")\n",
    "\n",
    "    if repo_dir.exists():\n",
    "        shutil.rmtree(repo_dir, onerror=on_rm_error)\n",
    "\n",
    "    try:\n",
    "        await asyncio.to_thread(git.Repo.clone_from, LICHESS_REPO_URL, str(repo_dir), depth=1)\n",
    "    except git.Exc as e:\n",
    "        logger.error(f\"Failed to clone Lichess repo: {e}\")\n",
    "        return\n",
    "\n",
    "    logger.info(\"Clone complete. Converting SVGs to PNGs...\")\n",
    "    source_path = repo_dir / target_subfolder\n",
    "\n",
    "    if not source_path.exists():\n",
    "        logger.error(f\"Could not find piece directory in repo: {source_path}\")\n",
    "        shutil.rmtree(repo_dir, onerror=on_rm_error)\n",
    "        return\n",
    "\n",
    "    svg_files = []\n",
    "    for root, _, files in os.walk(source_path):\n",
    "        root_path = Path(root)\n",
    "        if root_path == source_path:\n",
    "            continue\n",
    "\n",
    "        for file in files:\n",
    "            if file.lower().endswith(\".svg\"):\n",
    "                rel_dir = os.path.relpath(root, source_path)\n",
    "                output_dir_path = output_dir / rel_dir\n",
    "                svg_path = root_path / file\n",
    "                png_filename = os.path.splitext(file)[0] + \".png\"\n",
    "                png_path = output_dir_path / png_filename\n",
    "                \n",
    "                # Skip if already exists\n",
    "                if not png_path.exists():\n",
    "                    svg_files.append((svg_path, png_path))\n",
    "\n",
    "    logger.info(f\"Found {len(svg_files)} SVG files to convert\")\n",
    "\n",
    "    # Process in parallel using ProcessPoolExecutor\n",
    "    # We prefer ProcessPoolExecutor for CPU-bound SVG conversion\n",
    "    max_workers = max(1, os.cpu_count() - 1)\n",
    "    failed_count = 0\n",
    "    \n",
    "    if svg_files:\n",
    "        with ProcessPoolExecutor(max_workers=max_workers) as executor:\n",
    "            results = list(tqdm(\n",
    "                executor.map(convert_single_svg, svg_files),\n",
    "                total=len(svg_files),\n",
    "                desc=\"Converting Lichess SVGs (Parallel)\"\n",
    "            ))\n",
    "            failed_count = results.count(False)\n",
    "\n",
    "    if failed_count > 0:\n",
    "        logger.warning(f\"Note: {failed_count} SVG files failed to convert\")\n",
    "\n",
    "    logger.info(\"Lichess conversion complete. Cleaning up repo...\")\n",
    "    shutil.rmtree(repo_dir, onerror=on_rm_error)\n",
    "    logger.info(\"Cleanup complete.\")\n",
    "\n",
    "async def download_assets():\n",
    "    output_dir = DEFAULT_OUTPUT_DIR\n",
    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "    tasks = [\n",
    "        download_chesscom_pieces(output_dir),\n",
    "        download_lichess_pieces(output_dir)\n",
    "    ]\n",
    "    \n",
    "    await asyncio.gather(*tasks)\n",
    "    logger.info(f\"All assets downloaded to {output_dir.absolute()}\")\n",
    "\n",
    "# Run download\n",
    "await download_assets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 3: Data Generation (Parallelized)\n",
    "# This cell writes the generation logic to a script and runs it to ensure multiprocessing works correctly.\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# Define the script content\n",
    "data_gen_script = \"\"\"\n",
    "import os\n",
    "import sys\n",
    "import random\n",
    "import cv2\n",
    "import numpy as np\n",
    "import gc\n",
    "import shutil\n",
    "import time\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "from concurrent.futures import ProcessPoolExecutor, as_completed\n",
    "from tqdm import tqdm\n",
    "from typing import List, Tuple, Optional\n",
    "\n",
    "# --- Constants ---\n",
    "BOARD_THEMES = [\n",
    "    ((240, 217, 181), (181, 136, 99)),  # Classic brown\n",
    "    ((238, 238, 210), (118, 150, 86)),  # Green\n",
    "    ((222, 227, 230), (140, 162, 173)),  # Blue\n",
    "    ((255, 206, 158), (209, 139, 71)),  # Wood\n",
    "    ((235, 236, 208), (115, 149, 82)),  # Tournament\n",
    "    ((240, 240, 240), (120, 120, 120)),  # Gray\n",
    "    ((255, 228, 181), (205, 133, 63)),  # Tan\n",
    "    ((250, 235, 215), (139, 90, 43)),  # Marble\n",
    "    ((173, 216, 230), (100, 149, 237)),  # Sky blue\n",
    "    ((255, 218, 185), (160, 82, 45)),  # Peach\n",
    "    ((245, 222, 179), (139, 69, 19)),  # Wheat\n",
    "    ((230, 230, 250), (147, 112, 219)),  # Purple\n",
    "]\n",
    "\n",
    "# Global cache for workers\n",
    "PIECE_CACHE = {}\n",
    "\n",
    "# --- Core Logic ---\n",
    "\n",
    "def random_fen_generation() -> str:\n",
    "    pieces = [\"r\", \"n\", \"b\", \"q\", \"k\", \"p\", \"R\", \"N\", \"B\", \"Q\", \"K\", \"P\"]\n",
    "    fen = []\n",
    "    for _ in range(8):\n",
    "        row = [None] * 8\n",
    "        for i in range(8):\n",
    "            if random.random() < 0.5:\n",
    "                row[i] = random.choice(pieces)\n",
    "        row_str = \"\"\n",
    "        empty_count = 0\n",
    "        for piece in row:\n",
    "            if piece is not None:\n",
    "                if empty_count > 0:\n",
    "                    row_str += str(empty_count)\n",
    "                    empty_count = 0\n",
    "                row_str += piece\n",
    "            else:\n",
    "                empty_count += 1\n",
    "        if empty_count > 0:\n",
    "            row_str += str(empty_count)\n",
    "        fen.append(row_str)\n",
    "    return \"/\".join(fen) + \" w KQkq - 0 1\"\n",
    "\n",
    "def generate_empty_board_fen() -> str:\n",
    "    board_type = random.choice([\"completely_empty\", \"very_sparse\", \"sparse\", \"hard_negative_midgame\"])\n",
    "    if board_type == \"completely_empty\":\n",
    "        return \"8/8/8/8/8/8/8/8 w KQkq - 0 1\"\n",
    "    pieces = [\"r\", \"n\", \"b\", \"q\", \"k\", \"p\", \"R\", \"N\", \"B\", \"Q\", \"K\", \"P\"]\n",
    "    fen = []\n",
    "    piece_probability = 0.05 if board_type == \"very_sparse\" else 0.15 if board_type == \"sparse\" else 0.35\n",
    "    for _ in range(8):\n",
    "        row = [None] * 8\n",
    "        for i in range(8):\n",
    "            if random.random() < piece_probability:\n",
    "                row[i] = random.choice(pieces)\n",
    "        row_str = \"\"\n",
    "        empty_count = 0\n",
    "        for piece in row:\n",
    "            if piece is not None:\n",
    "                if empty_count > 0:\n",
    "                    row_str += str(empty_count)\n",
    "                    empty_count = 0\n",
    "                row_str += piece\n",
    "            else:\n",
    "                empty_count += 1\n",
    "        if empty_count > 0:\n",
    "            row_str += str(empty_count)\n",
    "        fen.append(row_str)\n",
    "    return \"/\".join(fen) + \" w KQkq - 0 1\"\n",
    "\n",
    "def generate_royalty_focused_fen() -> str:\n",
    "    fen = []\n",
    "    royalty_pieces = [\"k\", \"q\", \"b\", \"K\", \"Q\", \"B\"]\n",
    "    support_pieces = [\"r\", \"n\", \"p\", \"R\", \"N\", \"P\"]\n",
    "    for _ in range(8):\n",
    "        row = [None] * 8\n",
    "        for i in range(8):\n",
    "            if random.random() < 0.4:\n",
    "                if random.random() < 0.7:\n",
    "                    row[i] = random.choice(royalty_pieces)\n",
    "                else:\n",
    "                    row[i] = random.choice(support_pieces)\n",
    "        row_str = \"\"\n",
    "        empty_count = 0\n",
    "        for piece in row:\n",
    "            if piece is not None:\n",
    "                if empty_count > 0:\n",
    "                    row_str += str(empty_count)\n",
    "                    empty_count = 0\n",
    "                row_str += piece\n",
    "            else:\n",
    "                empty_count += 1\n",
    "        if empty_count > 0:\n",
    "            row_str += str(empty_count)\n",
    "        fen.append(row_str)\n",
    "    return \"/\".join(fen) + \" w KQkq - 0 1\"\n",
    "\n",
    "def fen_to_grid(fen: str) -> List[int]:\n",
    "    piece_to_id = {\"r\": 1, \"n\": 2, \"b\": 3, \"q\": 4, \"k\": 5, \"p\": 6, \"R\": 7, \"N\": 8, \"B\": 9, \"Q\": 10, \"K\": 11, \"P\": 12}\n",
    "    board_fen = fen.split()[0]\n",
    "    grid = []\n",
    "    for row in board_fen.split(\"/\"):\n",
    "        for char in row:\n",
    "            if char.isdigit():\n",
    "                grid.extend([0] * int(char))\n",
    "            else:\n",
    "                grid.append(piece_to_id[char])\n",
    "    return grid\n",
    "\n",
    "def fen_to_piece_list(fen: str) -> List[Tuple[str, int, int]]:\n",
    "    piece_map = {\"r\": \"bR\", \"n\": \"bN\", \"b\": \"bB\", \"q\": \"bQ\", \"k\": \"bK\", \"p\": \"bP\", \"R\": \"wR\", \"N\": \"wN\", \"B\": \"wB\", \"Q\": \"wQ\", \"K\": \"wK\", \"P\": \"wP\"}\n",
    "    board_fen = fen.split()[0]\n",
    "    pieces = []\n",
    "    for row_idx, row in enumerate(board_fen.split(\"/\")):\n",
    "        col_idx = 0\n",
    "        for char in row:\n",
    "            if char.isdigit():\n",
    "                col_idx += int(char)\n",
    "            else:\n",
    "                pieces.append((piece_map[char], row_idx, col_idx))\n",
    "                col_idx += 1\n",
    "    return pieces\n",
    "\n",
    "def get_available_piece_sets(assets_dir: Path) -> List[str]:\n",
    "    pieces_dir = assets_dir / \"pieces\"\n",
    "    if not pieces_dir.exists(): return []\n",
    "    blacklist = [\"blindfold\"]\n",
    "    piece_sets = []\n",
    "    for theme_dir in pieces_dir.iterdir():\n",
    "        if theme_dir.is_dir() and (theme_dir / \"wK.png\").exists():\n",
    "            if theme_dir.name.lower() not in blacklist:\n",
    "                piece_sets.append(theme_dir.name)\n",
    "    return sorted(piece_sets)\n",
    "\n",
    "def draw_chess_board(fen: str, piece_set: str, board_theme: Tuple, assets_dir: Path, board_size: int = 640) -> Optional[np.ndarray]:\n",
    "    global PIECE_CACHE\n",
    "    square_size = board_size // 8\n",
    "    light_color, dark_color = board_theme\n",
    "    \n",
    "    # Fast board creation using numpy\n",
    "    board = np.zeros((board_size, board_size, 3), dtype=np.uint8)\n",
    "    \n",
    "    # Fill board colors (optimized)\n",
    "    # Create 8x8 mask\n",
    "    y_indices, x_indices = np.indices((8, 8))\n",
    "    is_light = (y_indices + x_indices) % 2 == 0\n",
    "    \n",
    "    # Create small 8x8x3 grid\n",
    "    grid_colors = np.zeros((8, 8, 3), dtype=np.uint8)\n",
    "    grid_colors[is_light] = light_color\n",
    "    grid_colors[~is_light] = dark_color\n",
    "    \n",
    "    # Scale up to full board\n",
    "    board = grid_colors.repeat(square_size, axis=0).repeat(square_size, axis=1)\n",
    "\n",
    "    pieces = fen_to_piece_list(fen)\n",
    "    \n",
    "    # Load pieces for this set if not in cache (should be handled by worker_init, but fallback)\n",
    "    if piece_set not in PIECE_CACHE:\n",
    "        # Fallback load (slow)\n",
    "        piece_dir = assets_dir / \"pieces\" / piece_set\n",
    "        PIECE_CACHE[piece_set] = {}\n",
    "        for piece_name in [\"bR\", \"bN\", \"bB\", \"bQ\", \"bK\", \"bP\", \"wR\", \"wN\", \"wB\", \"wQ\", \"wK\", \"wP\"]:\n",
    "            p = piece_dir / f\"{piece_name}.png\"\n",
    "            if p.exists():\n",
    "                img = cv2.imread(str(p), cv2.IMREAD_UNCHANGED)\n",
    "                if img is not None:\n",
    "                     if img.shape[0] != square_size:\n",
    "                        img = cv2.resize(img, (square_size, square_size))\n",
    "                     PIECE_CACHE[piece_set][piece_name] = img\n",
    "    \n",
    "    piece_dict = PIECE_CACHE.get(piece_set, {})\n",
    "    \n",
    "    for piece_name, row, col in pieces:\n",
    "        piece_img = piece_dict.get(piece_name)\n",
    "        \n",
    "        if piece_img is None or len(piece_img.shape) < 3: continue\n",
    "        \n",
    "        y, x = row * square_size, col * square_size\n",
    "        \n",
    "        # Alpha blending\n",
    "        try:\n",
    "            if piece_img.shape[2] == 4:\n",
    "                # Optimization: Use pre-split channels if possible, but here we assume raw cv2 img\n",
    "                alpha = piece_img[:, :, 3:4] / 255.0\n",
    "                piece_rgb = piece_img[:, :, :3]\n",
    "                \n",
    "                roi = board[y : y + square_size, x : x + square_size]\n",
    "                \n",
    "                # Vectorized blending\n",
    "                blended = (alpha * piece_rgb + (1 - alpha) * roi)\n",
    "                board[y : y + square_size, x : x + square_size] = blended.astype(np.uint8)\n",
    "            else:\n",
    "                # No alpha, just overwrite\n",
    "                if len(piece_img.shape) == 2: \n",
    "                    piece_img = cv2.cvtColor(piece_img, cv2.COLOR_GRAY2BGR)\n",
    "                board[y : y + square_size, x : x + square_size] = piece_img[:, :, :3]\n",
    "        except Exception: \n",
    "            continue\n",
    "            \n",
    "    return board\n",
    "\n",
    "def add_gray_dots_to_board(board: np.ndarray, num_dots: Optional[int] = None) -> np.ndarray:\n",
    "    if board is None: return board\n",
    "    board_copy = board.copy()\n",
    "    square_size = board.shape[0] // 8\n",
    "    if num_dots is None: num_dots = random.randint(1, 5)\n",
    "    for _ in range(num_dots):\n",
    "        row, col = random.randint(0, 7), random.randint(0, 7)\n",
    "        center_y = int(row * square_size + square_size // 2)\n",
    "        center_x = int(col * square_size + square_size // 2)\n",
    "        dot_radius = int(random.randint(square_size // 8, square_size // 4))\n",
    "        gray_value = int(random.randint(120, 160))\n",
    "        dot_color = (gray_value, gray_value, gray_value)\n",
    "        offset_x = int(random.randint(-square_size // 6, square_size // 6))\n",
    "        offset_y = int(random.randint(-square_size // 6, square_size // 6))\n",
    "        try:\n",
    "            # Draw directly on copy\n",
    "            overlay = board_copy.copy()\n",
    "            cv2.circle(overlay, (center_x + offset_x, center_y + offset_y), dot_radius, dot_color, -1)\n",
    "            alpha = random.uniform(0.5, 0.8)\n",
    "            cv2.addWeighted(overlay, alpha, board_copy, 1 - alpha, 0, board_copy)\n",
    "        except Exception:\n",
    "            pass\n",
    "    return board_copy\n",
    "\n",
    "def generate_training_sample(assets_dir, piece_sets, output_size=640, is_negative=False, mode=\"normal\") -> Tuple[Optional[np.ndarray], str]:\n",
    "    if mode == \"hard_negative_dots\": fen = generate_empty_board_fen()\n",
    "    elif mode == \"royalty_focus\": fen = generate_royalty_focused_fen()\n",
    "    elif is_negative: fen = generate_empty_board_fen()\n",
    "    else: fen = random_fen_generation()\n",
    "    \n",
    "    piece_set = random.choice(piece_sets)\n",
    "    board_theme = random.choice(BOARD_THEMES)\n",
    "    \n",
    "    board = draw_chess_board(fen, piece_set, board_theme, assets_dir, output_size)\n",
    "    if board is None: return None, fen\n",
    "    \n",
    "    if mode == \"hard_negative_dots\": \n",
    "        board = add_gray_dots_to_board(board)\n",
    "        \n",
    "    # Optional: Add noise/brightness (removed perspective/scale here as requested)\n",
    "    if random.random() < 0.3:\n",
    "        try:\n",
    "            alpha, beta = random.uniform(0.9, 1.1), random.uniform(-5, 5)\n",
    "            board = cv2.convertScaleAbs(board, alpha=alpha, beta=beta)\n",
    "        except Exception:\n",
    "            pass\n",
    "            \n",
    "    return board, fen\n",
    "\n",
    "# --- Worker Function ---\n",
    "\n",
    "def worker_init(assets_dir_str, piece_sets, square_size):\n",
    "    # Pre-load assets for this worker\n",
    "    global PIECE_CACHE\n",
    "    PIECE_CACHE.clear()\n",
    "    assets_dir = Path(assets_dir_str)\n",
    "    \n",
    "    # Optimization: Only load assigned piece sets\n",
    "    for piece_set in piece_sets:\n",
    "        piece_dir = assets_dir / \"pieces\" / piece_set\n",
    "        if not piece_dir.exists(): continue\n",
    "        \n",
    "        PIECE_CACHE[piece_set] = {}\n",
    "        for piece_name in [\"bR\", \"bN\", \"bB\", \"bQ\", \"bK\", \"bP\", \"wR\", \"wN\", \"wB\", \"wQ\", \"wK\", \"wP\"]:\n",
    "            p = piece_dir / f\"{piece_name}.png\"\n",
    "            if p.exists():\n",
    "                img = cv2.imread(str(p), cv2.IMREAD_UNCHANGED)\n",
    "                if img is not None:\n",
    "                    # Resize once here\n",
    "                    if img.shape[0] != square_size:\n",
    "                        img = cv2.resize(img, (square_size, square_size), interpolation=cv2.INTER_AREA)\n",
    "                    PIECE_CACHE[piece_set][piece_name] = img\n",
    "\n",
    "def generate_batch(start_idx, n_samples, split_name, piece_sets, assets_dir_str, output_dir_str, mode, seed, save_clean_boards, boards_dir_str):\n",
    "    # Initialize worker logic (re-seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed % (2**32))\n",
    "    \n",
    "    assets_dir = Path(assets_dir_str)\n",
    "    output_dir = Path(output_dir_str)\n",
    "    boards_dir = Path(boards_dir_str) if boards_dir_str else None\n",
    "    \n",
    "    # Ensure caching (if worker_init wasn't called or lost context)\n",
    "    global PIECE_CACHE\n",
    "    if not PIECE_CACHE:\n",
    "        worker_init(assets_dir_str, piece_sets, 640//8)\n",
    "        \n",
    "    images_dir = output_dir / split_name / \"images\"\n",
    "    labels_dir = output_dir / split_name / \"labels\"\n",
    "    \n",
    "    processed = 0\n",
    "    for i in range(n_samples):\n",
    "        idx = start_idx + i\n",
    "        \n",
    "        # Generate\n",
    "        board, fen = generate_training_sample(assets_dir, piece_sets, output_size=640, is_negative=(mode==\"neg\"), mode=mode)\n",
    "        \n",
    "        if board is None: continue\n",
    "        \n",
    "        try:\n",
    "            # Save Image (Use PIL for speed/compat)\n",
    "            # BGR to RGB\n",
    "            img_rgb = board[:, :, ::-1]\n",
    "            \n",
    "            img_path = images_dir / f\"sample_{idx:06d}.png\"\n",
    "            Image.fromarray(img_rgb).save(str(img_path))\n",
    "            \n",
    "            # Save Labels\n",
    "            grid = fen_to_grid(fen)\n",
    "            labels = []\n",
    "            for idx2, piece_id in enumerate(grid):\n",
    "                if piece_id == 0: continue\n",
    "                row, col = idx2 // 8, idx2 % 8\n",
    "                center_x, center_y = (col + 0.5) / 8, (row + 0.5) / 8\n",
    "                labels.append(f\"{piece_id - 1} {center_x:.6f} {center_y:.6f} {1.0/8:.6f} {1.0/8:.6f}\")\n",
    "            \n",
    "            lbl_path = labels_dir / f\"sample_{idx:06d}.txt\"\n",
    "            with open(lbl_path, \"w\") as f:\n",
    "                f.write(\"\\\\n\".join(labels) + \"\\\\n\")\n",
    "                \n",
    "            # Save clean board if requested\n",
    "            if save_clean_boards and boards_dir:\n",
    "                Image.fromarray(img_rgb).save(str(boards_dir / f\"{split_name}_{idx:06d}.png\"))\n",
    "                \n",
    "            processed += 1\n",
    "        except Exception as e:\n",
    "            pass\n",
    "            \n",
    "    return processed\n",
    "\n",
    "# --- Orchestration ---\n",
    "\n",
    "def generate_dataset_parallel(assets_dir, output_dir, count, negative_ratio, hard_negative_dots_count, royalty_focus_count, save_clean_boards):\n",
    "    assets_path = Path(assets_dir)\n",
    "    output_path = Path(output_dir)\n",
    "    \n",
    "    # Setup Dirs\n",
    "    if output_path.exists(): shutil.rmtree(output_path)\n",
    "    for split in [\"train\", \"val\"]:\n",
    "        (output_path / split / \"images\").mkdir(parents=True, exist_ok=True)\n",
    "        (output_path / split / \"labels\").mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "    boards_path = None\n",
    "    if save_clean_boards:\n",
    "        boards_path = assets_path / \"boards\"\n",
    "        if boards_path.exists(): shutil.rmtree(boards_path)\n",
    "        boards_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    # Get Piece Sets\n",
    "    piece_sets = get_available_piece_sets(assets_path)\n",
    "    if not piece_sets:\n",
    "        print(f\"ERROR: No piece sets found in {assets_path}/pieces!\")\n",
    "        print(\"Please check your asset download.\")\n",
    "        return\n",
    "    \n",
    "    print(f\"Found {len(piece_sets)} piece sets. Starting generation...\")\n",
    "    random.shuffle(piece_sets)\n",
    "    \n",
    "    # Splits\n",
    "    split_idx = max(1, int(len(piece_sets) * 0.8))\n",
    "    train_sets, val_sets = piece_sets[:split_idx], piece_sets[split_idx:]\n",
    "    \n",
    "    # Calc Counts\n",
    "    total_samples = int(count / (1 - negative_ratio))\n",
    "    neg_count = int(total_samples * negative_ratio)\n",
    "    pos_count = count\n",
    "    \n",
    "    # Task Configs\n",
    "    tasks = []\n",
    "    \n",
    "    def add_task_group(split, sets, mode, n_total):\n",
    "        # Split n_total into batches of ~1000\n",
    "        batch_size = 1000\n",
    "        n_batches = (n_total + batch_size - 1) // batch_size\n",
    "        \n",
    "        base_idx = len(tasks) * 100000 # Offset to avoid index collision across types? No, use sequential\n",
    "        # Better: pass start_idx dynamically\n",
    "        return n_batches, batch_size\n",
    "\n",
    "    # We'll just create a list of work items\n",
    "    # Work Item: (start_idx, count, split, mode, sets)\n",
    "    \n",
    "    work_items = []\n",
    "    global_idx = 0\n",
    "    \n",
    "    # Helper to schedule\n",
    "    def schedule(n, split, sets, mode):\n",
    "        nonlocal global_idx\n",
    "        batch_size = 1000\n",
    "        remaining = n\n",
    "        while remaining > 0:\n",
    "            b = min(batch_size, remaining)\n",
    "            work_items.append({\n",
    "                \"start_idx\": global_idx,\n",
    "                \"n\": b,\n",
    "                \"split\": split,\n",
    "                \"sets\": sets,\n",
    "                \"mode\": mode,\n",
    "                \"seed\": random.randint(0, 1000000)\n",
    "            })\n",
    "            global_idx += b\n",
    "            remaining -= b\n",
    "            \n",
    "    # Train\n",
    "    schedule(int(pos_count * 0.8), \"train\", train_sets, \"normal\")\n",
    "    schedule(int(neg_count * 0.8), \"train\", train_sets, \"neg\")\n",
    "    schedule(int(hard_negative_dots_count * 0.8), \"train\", train_sets, \"hard_negative_dots\")\n",
    "    schedule(int(royalty_focus_count * 0.8), \"train\", train_sets, \"royalty_focus\")\n",
    "    \n",
    "    # Val\n",
    "    schedule(pos_count - int(pos_count * 0.8), \"val\", val_sets, \"normal\")\n",
    "    schedule(neg_count - int(neg_count * 0.8), \"val\", val_sets, \"neg\")\n",
    "    schedule(hard_negative_dots_count - int(hard_negative_dots_count * 0.8), \"val\", val_sets, \"hard_negative_dots\")\n",
    "    schedule(royalty_focus_count - int(royalty_focus_count * 0.8), \"val\", val_sets, \"royalty_focus\")\n",
    "    \n",
    "    print(f\"Scheduled {len(work_items)} batches. Total samples: {global_idx}\")\n",
    "    \n",
    "    # Run Parallel\n",
    "    # Max workers = CPU count\n",
    "    max_workers = os.cpu_count() or 4\n",
    "    print(f\"Using {max_workers} workers.\")\n",
    "    \n",
    "    with ProcessPoolExecutor(max_workers=max_workers) as executor:\n",
    "        futures = []\n",
    "        for item in work_items:\n",
    "            f = executor.submit(\n",
    "                generate_batch,\n",
    "                item[\"start_idx\"],\n",
    "                item[\"n\"],\n",
    "                item[\"split\"],\n",
    "                item[\"sets\"],\n",
    "                str(assets_path),\n",
    "                str(output_path),\n",
    "                item[\"mode\"],\n",
    "                item[\"seed\"],\n",
    "                save_clean_boards,\n",
    "                str(boards_path) if boards_path else None\n",
    "            )\n",
    "            futures.append(f)\n",
    "            \n",
    "        # Progress\n",
    "        total_done = 0\n",
    "        for f in tqdm(as_completed(futures), total=len(futures), desc=\"Generating\"):\n",
    "            total_done += f.result()\n",
    "            \n",
    "    print(f\"Generation complete. {total_done} images generated.\")\n",
    "    \n",
    "    # Create YAML\n",
    "    yaml_content = f\"path: {output_path.absolute().as_posix()}\\\\ntrain: train/images\\\\nval: val/images\\\\nnames:\\\\n\"\n",
    "    for i, n in enumerate(['r','n','b','q','k','p','R','N','B','Q','K','P']): yaml_content += f\"  {i}: {n}\\\\n\"\n",
    "    with open(output_path / \"data.yaml\", \"w\") as f: f.write(yaml_content)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    generate_dataset_parallel(\n",
    "        assets_dir=\"./assets\",\n",
    "        output_dir=\"yolo_pieces\",\n",
    "        count=15000,\n",
    "        negative_ratio=0.15,\n",
    "        hard_negative_dots_count=2000,\n",
    "        royalty_focus_count=3000,\n",
    "        save_clean_boards=True\n",
    "    )\n",
    "\"\"\"\n",
    "\n",
    "# Write to file\n",
    "with open(\"data_gen.py\", \"w\") as f:\n",
    "    f.write(data_gen_script)\n",
    "\n",
    "# Run the script\n",
    "print(\"Running data generation script...\")\n",
    "!python data_gen.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4: Training\n",
    "!pip install \"numpy<2\"\n",
    "\n",
    "from ultralytics import YOLO\n",
    "from IPython.display import Image, display\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "def train_model(config_yaml, project_name, base_model='yolov8n.pt', epochs=50, imgsz=640, device_ids=[0, 1]):\n",
    "    try:\n",
    "        model = YOLO(base_model) \n",
    "    except Exception as e:\n",
    "        print(f\"Error loading base model {base_model}: {e}\")\n",
    "        return None\n",
    "        \n",
    "    print(f\"\\n--- Starting Training for {project_name} using {base_model} ---\")\n",
    "    \n",
    "    # Handle devices list for single GPU or CPU\n",
    "    if isinstance(device_ids, int):\n",
    "        device_ids = [device_ids]\n",
    "    elif device_ids == 'cpu':\n",
    "        device_ids = 'cpu'\n",
    "\n",
    "    print(f\"Using device: {device_ids}\")\n",
    "    \n",
    "    # Optimize workers based on CPU count\n",
    "    workers = min(os.cpu_count(), 16)\n",
    "    print(f\"Using {workers} workers for data loading\")\n",
    "\n",
    "    try:\n",
    "        # Train the model\n",
    "        model.train(\n",
    "            data=os.path.join(WORKING_DIR, config_yaml), \n",
    "            epochs=epochs, \n",
    "            imgsz=imgsz, \n",
    "            device=device_ids,\n",
    "            patience=20,\n",
    "            batch=32, \n",
    "            workers=workers, \n",
    "            # Augmentations settings for \"Screen/Screenshot\" use case\n",
    "            # Disable geometric distortions\n",
    "            augment=True, # Enables other augs like HSV, Mosaic\n",
    "            perspective=0.0, \n",
    "            degrees=0.0,    # No rotation\n",
    "            translate=0.1,  # Small translation is OK (crop position variation)\n",
    "            scale=0.0,      # No scaling\n",
    "            shear=0.0,\n",
    "            mosaic=0.0,     # Disable mosaic to preserve rigid grid structure\n",
    "            mixup=0.0,\n",
    "            flipud=0.0,     # No vertical flip\n",
    "            fliplr=0.0,     # No horizontal flip (chess board is asymmetric)\n",
    "            # Color/Lighting augs are still good for different screens/themes\n",
    "            hsv_v=0.4, \n",
    "            hsv_s=0.7,\n",
    "            \n",
    "            project=project_name,\n",
    "            name='train',\n",
    "            cache=False,\n",
    "            exist_ok=False \n",
    "        )\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred during training: {e}\")\n",
    "        return None\n",
    "    \n",
    "    # Retrieve the actual save directory from the model trainer\n",
    "    if hasattr(model, 'trainer') and hasattr(model.trainer, 'save_dir'):\n",
    "        save_dir = model.trainer.save_dir\n",
    "        best_model_path = os.path.join(save_dir, 'weights', 'best.pt')\n",
    "        print(f\"Model saved to {save_dir}\")\n",
    "    else:\n",
    "        print(\"Warning: Could not retrieve save_dir from model.trainer. Checking default path.\")\n",
    "        best_model_path = os.path.join(project_name, 'train', 'weights', 'best.pt')\n",
    "\n",
    "    final_path = os.path.join(WORKING_DIR, f'{project_name}_best.pt')\n",
    "    \n",
    "    if os.path.exists(best_model_path):\n",
    "        shutil.copy(best_model_path, final_path)\n",
    "        print(f\"Saved best model to {final_path}\")\n",
    "        try:\n",
    "            results_png = os.path.join(os.path.dirname(os.path.dirname(best_model_path)), 'results.png')\n",
    "            if os.path.exists(results_png):\n",
    "                display(Image(filename=results_png))\n",
    "        except Exception:\n",
    "            pass\n",
    "        return final_path\n",
    "    \n",
    "    print(f\"Best model weights not found at {best_model_path}\")\n",
    "    return None\n",
    "\n",
    "# Determine device\n",
    "if torch.cuda.is_available():\n",
    "    device_config = [i for i in range(torch.cuda.device_count())]\n",
    "else:\n",
    "    device_config = 'cpu'\n",
    "\n",
    "# Start training\n",
    "MODEL_PATH = train_model(\n",
    "    'yolo_pieces/data.yaml', \n",
    "    'PieceDetector_Small', \n",
    "    base_model='yolov8s.pt', \n",
    "    epochs=150, \n",
    "    imgsz=640,\n",
    "    device_ids=device_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5: Export to ONNX\n",
    "def convert_model_to_onnx(model_path: Path, output_path: Optional[Path] = None) -> bool:\n",
    "    logger.info(f\"Converting {model_path}...\")\n",
    "    try:\n",
    "        model = YOLO(str(model_path))\n",
    "        onnx_file = model.export(format=\"onnx\", simplify=True)\n",
    "        logger.info(f\"Successfully converted to {onnx_file}\")\n",
    "        if output_path:\n",
    "             generated_path = Path(onnx_file)\n",
    "             if generated_path.resolve() != output_path.resolve():\n",
    "                 shutil.move(generated_path, output_path)\n",
    "                 logger.info(f\"Moved to {output_path}\")\n",
    "        return True\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Failed to convert {model_path}: {e}\")\n",
    "        return False\n",
    "\n",
    "if MODEL_PATH and os.path.exists(MODEL_PATH):\n",
    "    print(f\"Converting trained model: {MODEL_PATH}\")\n",
    "    convert_model_to_onnx(Path(MODEL_PATH), Path(MODEL_PATH).with_suffix('.onnx'))\n",
    "else:\n",
    "    print(\"No model found to export.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
