train_cursor.py

"""
End-to-end cursor detection: generate synthetic YOLO dataset + train YOLOv8n.
"""
import os
import random
from pathlib import Path
from PIL import Image
from datasets import load_dataset
from tqdm import tqdm
from ultralytics import YOLO

# ========== CONFIG ==========
NUM_TRAIN = 500
NUM_VAL = 100
NUM_TEST = 50
IMAGE_SIZE = (640, 640)
CURSOR_SIZE_RANGE = (16, 48)
OUTPUT_DIR = Path("/app/cursor_dataset")
YAML_PATH = OUTPUT_DIR / "cursor.yaml"

# ========== SETUP ==========
for split in ["train", "val", "test"]:
    (OUTPUT_DIR / "images" / split).mkdir(parents=True, exist_ok=True)
    (OUTPUT_DIR / "labels" / split).mkdir(parents=True, exist_ok=True)

print("Loading datasets...")
cursors_ds = load_dataset("Fraser/cursors", split="train")
screenshots_ds = load_dataset("naorm/website-screenshots", split="train")

static_cursors = []
animated_cursors = []
for row in cursors_ds:
    frames = row["frames"]
    if len(frames) == 1:
        static_cursors.append(row)
    else:
        animated_cursors.append(row)

print(f"Cursors: {len(static_cursors)} static, {len(animated_cursors)} animated")
print(f"Screenshots: {len(screenshots_ds)}")

screenshot_images = [row["image"] for row in screenshots_ds]


def get_random_cursor():
    if random.random() < 0.7 and static_cursors:
        row = random.choice(static_cursors)
        frame = row["frames"][0]
    elif animated_cursors:
        row = random.choice(animated_cursors)
        frame = random.choice(row["frames"])
    else:
        row = random.choice(static_cursors)
        frame = row["frames"][0]
    return frame, row["hotspot_x"], row["hotspot_y"], row["size_px"]


def composite_cursor(bg_img, cursor_frame, hotspot_x, hotspot_y, cursor_size):
    bg = bg_img.convert("RGBA").resize(IMAGE_SIZE, Image.LANCZOS)
    cw, ch = cursor_frame.size
    scale = cursor_size / max(cw, ch)
    new_w = max(1, round(cw * scale))
    new_h = max(1, round(ch * scale))
    cursor_scaled = cursor_frame.convert("RGBA").resize((new_w, new_h), Image.LANCZOS)

    target_x = random.randint(0, IMAGE_SIZE[0] - 1)
    target_y = random.randint(0, IMAGE_SIZE[1] - 1)

    paste_x = target_x - round(hotspot_x * new_w)
    paste_y = target_y - round(hotspot_y * new_h)

    sw, sh = bg.size
    src_left = max(0, -paste_x)
    src_top = max(0, -paste_y)
    src_right = min(new_w, sw - paste_x)
    src_bottom = min(new_h, sh - paste_y)

    if src_left >= src_right or src_top >= src_bottom:
        return None, None

    cropped = cursor_scaled.crop((src_left, src_top, src_right, src_bottom))
    dst_x = paste_x + src_left
    dst_y = paste_y + src_top
    bg.alpha_composite(cropped, (dst_x, dst_y))

    bbox_x = (dst_x + (src_right - src_left) / 2) / sw
    bbox_y = (dst_y + (src_bottom - src_top) / 2) / sh
    bbox_w = (src_right - src_left) / sw
    bbox_h = (src_bottom - src_top) / sh

    bbox_x = max(bbox_w / 2, min(1 - bbox_w / 2, bbox_x))
    bbox_y = max(bbox_h / 2, min(1 - bbox_h / 2, bbox_y))
    bbox_w = min(1.0, bbox_w)
    bbox_h = min(1.0, bbox_h)

    return bg.convert("RGB"), (0, bbox_x, bbox_y, bbox_w, bbox_h)


def generate_split(split_name, num_samples):
    print(f"Generating {split_name} ({num_samples} samples)...")
    img_dir = OUTPUT_DIR / "images" / split_name
    lbl_dir = OUTPUT_DIR / "labels" / split_name

    for i in tqdm(range(num_samples), desc=split_name):
        bg = random.choice(screenshot_images)
        cursor_frame, hs_x, hs_y, size_px = get_random_cursor()
        cursor_size = random.randint(*CURSOR_SIZE_RANGE)

        composed, bbox = composite_cursor(bg, cursor_frame, hs_x, hs_y, cursor_size)

        retries = 0
        while composed is None and retries < 5:
            bg = random.choice(screenshot_images)
            cursor_frame, hs_x, hs_y, size_px = get_random_cursor()
            cursor_size = random.randint(*CURSOR_SIZE_RANGE)
            composed, bbox = composite_cursor(bg, cursor_frame, hs_x, hs_y, cursor_size)
            retries += 1

        if composed is None:
            continue

        img_path = img_dir / f"{split_name}_{i:06d}.jpg"
        lbl_path = lbl_dir / f"{split_name}_{i:06d}.txt"

        composed.save(img_path, "JPEG", quality=95)
        with open(lbl_path, "w") as f:
            f.write(f"0 {bbox[1]:.6f} {bbox[2]:.6f} {bbox[3]:.6f} {bbox[4]:.6f}\n")


random.seed(42)
generate_split("train", NUM_TRAIN)
generate_split("val", NUM_VAL)
generate_split("test", NUM_TEST)

# Write YAML
with open(YAML_PATH, "w") as f:
    f.write(f"path: {OUTPUT_DIR.absolute()}\n")
    f.write("train: images/train\n")
    f.write("val: images/val\n")
    f.write("test: images/test\n")
    f.write("nc: 1\n")
    f.write("names: ['cursor']\n")

print(f"Dataset ready at {OUTPUT_DIR}")

# ========== TRAIN YOLO ==========
model = YOLO("yolov8n.pt")

results = model.train(
    data=str(YAML_PATH),
    epochs=30,
    imgsz=640,
    batch=16,
    device=0,
    project="cursor-detection",
    name="yolov8n-cursor",
    exist_ok=True,
    verbose=True,
)

# Evaluate on test set
metrics = model.val(data=str(YAML_PATH), split="test")
print("Test metrics:", metrics)

# Push to Hub
model.push_to_hub("AdithyaSK/cursor-detection-yolov8n")
print("Model pushed to Hub!")