training/prepare_surface_gate_dataset.py

from __future__ import annotations

import argparse
import json
import random
import zipfile
from pathlib import Path

import cv2
import numpy as np
import yaml
from huggingface_hub import snapshot_download
from tqdm import tqdm
from torchvision.datasets import CIFAR10, STL10


SEVERSTAL_REPO = "rohanath/severstal-steel-detection"
NEU_REPO = "LiuErXiao/NEU_valid"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Prepare a lightweight steel-surface detector dataset with automatic negatives and synthetic composites."
    )
    parser.add_argument(
        "--output",
        default="training/data/surface_gate_detector",
        help="Output folder for the YOLO dataset.",
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=640,
        help="Target square size for generated detector images.",
    )
    parser.add_argument(
        "--max-severstal",
        type=int,
        default=900,
        help="Maximum Severstal steel images to use.",
    )
    parser.add_argument(
        "--max-neu",
        type=int,
        default=400,
        help="Maximum NEU steel images to use.",
    )
    parser.add_argument(
        "--include-severstal",
        action="store_true",
        help="Download and include the larger Severstal source for a stronger but slower dataset build.",
    )
    parser.add_argument(
        "--max-negatives",
        type=int,
        default=1200,
        help="Maximum Imagenette images to use as automatic non-steel backgrounds.",
    )
    parser.add_argument(
        "--composite-multiplier",
        type=float,
        default=1.5,
        help="How many synthetic positive composites to generate per steel image.",
    )
    parser.add_argument(
        "--negative-multiplier",
        type=float,
        default=1.0,
        help="How many pure background negatives to generate per steel image.",
    )
    parser.add_argument(
        "--negative-source",
        default="cifar10",
        choices=["cifar10", "stl10", "auto"],
        help="Automatic non-steel image source. `cifar10` is the fastest default.",
    )
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args()


def ensure_clean_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def list_images(path: Path) -> list[Path]:
    suffixes = {".jpg", ".jpeg", ".png", ".bmp"}
    return [file for file in path.rglob("*") if file.suffix.lower() in suffixes]


def extract_first_zip(repo_id: str, target_dir: Path) -> Path:
    target_dir.mkdir(parents=True, exist_ok=True)
    downloaded_dir = Path(
        snapshot_download(
            repo_id=repo_id,
            repo_type="dataset",
            allow_patterns=["*.zip"],
        )
    )
    zip_files = sorted(downloaded_dir.rglob("*.zip"))
    if not zip_files:
        raise FileNotFoundError(f"No zip file found in dataset repo {repo_id}")

    marker = target_dir / ".extracted"
    if marker.exists():
        return target_dir

    with zipfile.ZipFile(zip_files[0]) as archive:
        archive.extractall(target_dir)
    marker.write_text("ok", encoding="utf-8")
    return target_dir


def sample_paths(paths: list[Path], limit: int, rng: random.Random) -> list[Path]:
    if len(paths) <= limit:
        return list(paths)
    return rng.sample(paths, limit)


def center_crop_and_resize(image: np.ndarray, size: int) -> np.ndarray:
    height, width = image.shape[:2]
    crop_size = min(height, width)
    x0 = max(0, (width - crop_size) // 2)
    y0 = max(0, (height - crop_size) // 2)
    cropped = image[y0:y0 + crop_size, x0:x0 + crop_size]
    return cv2.resize(cropped, (size, size), interpolation=cv2.INTER_AREA)


def random_steel_crop(image: np.ndarray, rng: random.Random) -> np.ndarray:
    height, width = image.shape[:2]
    crop_w = max(64, int(width * rng.uniform(0.35, 0.9)))
    crop_h = max(64, int(height * rng.uniform(0.45, 0.95)))
    x0 = rng.randint(0, max(width - crop_w, 0))
    y0 = rng.randint(0, max(height - crop_h, 0))
    return image[y0:y0 + crop_h, x0:x0 + crop_w]


def make_composite(
    steel_image: np.ndarray,
    background_image: np.ndarray,
    size: int,
    rng: random.Random,
) -> tuple[np.ndarray, tuple[int, int, int, int]]:
    canvas = center_crop_and_resize(background_image, size)
    steel_crop = random_steel_crop(steel_image, rng)

    target_w = int(size * rng.uniform(0.45, 0.92))
    aspect_ratio = steel_crop.shape[0] / max(steel_crop.shape[1], 1)
    target_h = int(target_w * aspect_ratio)
    target_h = max(int(size * 0.18), min(target_h, int(size * 0.82)))

    steel_patch = cv2.resize(steel_crop, (target_w, target_h), interpolation=cv2.INTER_AREA)
    x0 = rng.randint(0, max(size - target_w, 0))
    y0 = rng.randint(0, max(size - target_h, 0))

    alpha = np.ones((target_h, target_w), dtype=np.float32)
    alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=5, sigmaY=5)
    alpha = np.clip(alpha[..., None], 0.86, 1.0)

    roi = canvas[y0:y0 + target_h, x0:x0 + target_w].astype(np.float32)
    patch = steel_patch.astype(np.float32)
    mixed = cv2.convertScaleAbs((patch * alpha) + (roi * (1.0 - alpha)))
    canvas[y0:y0 + target_h, x0:x0 + target_w] = mixed

    return canvas, (x0, y0, target_w, target_h)


def write_yolo_label(label_path: Path, bbox: tuple[int, int, int, int] | None, image_size: int) -> None:
    if bbox is None:
        label_path.write_text("", encoding="utf-8")
        return

    x, y, w, h = bbox
    x_center = (x + (w / 2)) / image_size
    y_center = (y + (h / 2)) / image_size
    width_norm = w / image_size
    height_norm = h / image_size
    label_path.write_text(
        f"0 {x_center:.6f} {y_center:.6f} {width_norm:.6f} {height_norm:.6f}\n",
        encoding="utf-8",
    )


def save_image(path: Path, image: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(path), image)


def export_negative_pool(
    target_dir: Path,
    limit: int,
    rng: random.Random,
    negative_source: str,
) -> tuple[list[Path], str]:
    source_loaders = {
        "cifar10": lambda root: CIFAR10(root=str(root), train=True, download=True),
        "stl10": lambda root: STL10(root=str(root), split="train", download=True),
    }
    if negative_source == "auto":
        datasets_to_try = [("cifar10", source_loaders["cifar10"]), ("stl10", source_loaders["stl10"])]
    else:
        datasets_to_try = [(negative_source, source_loaders[negative_source])]

    last_error: Exception | None = None
    for dataset_name, loader in datasets_to_try:
        try:
            dataset = loader(target_dir / "_torchvision_cache")
            samples = list(range(len(dataset)))
            rng.shuffle(samples)

            output_paths: list[Path] = []
            for index, sample_index in enumerate(samples[:limit]):
                sample_image, _ = dataset[sample_index]
                image = cv2.cvtColor(np.array(sample_image), cv2.COLOR_RGB2BGR)
                destination = target_dir / f"{dataset_name}_{index:05d}.jpg"
                save_image(destination, center_crop_and_resize(image, 640))
                output_paths.append(destination)

            if output_paths:
                return output_paths, dataset_name
        except Exception as exc:  # pragma: no cover - download/runtime fallback
            last_error = exc

    raise RuntimeError(
        "Unable to download an automatic negative-image source via torchvision."
    ) from last_error


def build_split(
    split_name: str,
    steel_paths: list[Path],
    negative_paths: list[Path],
    images_dir: Path,
    labels_dir: Path,
    image_size: int,
    composite_multiplier: float,
    negative_multiplier: float,
    rng: random.Random,
) -> dict[str, int]:
    counts = {"full_positive": 0, "synthetic_positive": 0, "background_negative": 0}
    split_images = images_dir / split_name
    split_labels = labels_dir / split_name
    ensure_clean_dir(split_images)
    ensure_clean_dir(split_labels)

    for index, steel_path in enumerate(tqdm(steel_paths, desc=f"{split_name}: full-frame steel")):
        image = cv2.imread(str(steel_path))
        if image is None:
            continue

        output_image = center_crop_and_resize(image, image_size)
        image_path = split_images / f"{split_name}_steel_{index:05d}.jpg"
        label_path = split_labels / f"{split_name}_steel_{index:05d}.txt"
        save_image(image_path, output_image)
        margin = int(image_size * 0.02)
        write_yolo_label(
            label_path,
            (margin, margin, image_size - (margin * 2), image_size - (margin * 2)),
            image_size,
        )
        counts["full_positive"] += 1

    synthetic_target = max(1, int(len(steel_paths) * composite_multiplier))
    for index in tqdm(range(synthetic_target), desc=f"{split_name}: synthetic composites"):
        steel_image = cv2.imread(str(rng.choice(steel_paths)))
        background_image = cv2.imread(str(rng.choice(negative_paths)))
        if steel_image is None or background_image is None:
            continue
        composite, bbox = make_composite(steel_image, background_image, image_size, rng)
        image_path = split_images / f"{split_name}_composite_{index:05d}.jpg"
        label_path = split_labels / f"{split_name}_composite_{index:05d}.txt"
        save_image(image_path, composite)
        write_yolo_label(label_path, bbox, image_size)
        counts["synthetic_positive"] += 1

    negative_target = max(1, int(len(steel_paths) * negative_multiplier))
    negative_sample = [rng.choice(negative_paths) for _ in range(negative_target)]
    for index, negative_path in enumerate(tqdm(negative_sample, desc=f"{split_name}: negative backgrounds")):
        image = cv2.imread(str(negative_path))
        if image is None:
            continue
        output_image = center_crop_and_resize(image, image_size)
        image_path = split_images / f"{split_name}_negative_{index:05d}.jpg"
        label_path = split_labels / f"{split_name}_negative_{index:05d}.txt"
        save_image(image_path, output_image)
        write_yolo_label(label_path, None, image_size)
        counts["background_negative"] += 1

    return counts


def write_dataset_yaml(dataset_root: Path) -> Path:
    yaml_path = dataset_root / "dataset.yaml"
    payload = {
        "path": str(dataset_root.resolve()),
        "train": "images/train",
        "val": "images/val",
        "test": "images/test",
        "names": {0: "steel_surface"},
    }
    yaml_path.write_text(yaml.safe_dump(payload, sort_keys=False), encoding="utf-8")
    return yaml_path


def main() -> None:
    args = parse_args()
    rng = random.Random(args.seed)

    output_root = Path(args.output).resolve()
    raw_root = output_root / "raw"
    dataset_root = output_root / "yolo_dataset"
    images_dir = dataset_root / "images"
    labels_dir = dataset_root / "labels"

    ensure_clean_dir(raw_root)
    ensure_clean_dir(images_dir)
    ensure_clean_dir(labels_dir)

    neu_raw = extract_first_zip(NEU_REPO, raw_root / "neu")
    negative_dir = raw_root / "negatives"
    ensure_clean_dir(negative_dir)

    steel_paths: list[Path] = []
    if args.include_severstal:
        severstal_raw = extract_first_zip(SEVERSTAL_REPO, raw_root / "severstal")
        steel_paths.extend(sample_paths(list_images(severstal_raw), args.max_severstal, rng))

    steel_paths.extend(sample_paths(list_images(neu_raw), args.max_neu, rng))
    steel_paths.extend(sorted(Path("test_images").glob("*.jpg")))
    steel_paths = [path for path in steel_paths if path.exists()]
    rng.shuffle(steel_paths)

    if not steel_paths:
        raise RuntimeError("No steel images were collected for the detector dataset.")

    negative_paths, negative_source = export_negative_pool(
        negative_dir,
        args.max_negatives,
        rng,
        args.negative_source,
    )
    if not negative_paths:
        raise RuntimeError("No negative background images were collected.")

    total = len(steel_paths)
    train_end = int(total * 0.8)
    val_end = int(total * 0.9)
    splits = {
        "train": steel_paths[:train_end],
        "val": steel_paths[train_end:val_end],
        "test": steel_paths[val_end:],
    }

    summary = {}
    for split_name, split_steel_paths in splits.items():
        summary[split_name] = build_split(
            split_name=split_name,
            steel_paths=split_steel_paths,
            negative_paths=negative_paths,
            images_dir=images_dir,
            labels_dir=labels_dir,
            image_size=args.image_size,
            composite_multiplier=args.composite_multiplier,
            negative_multiplier=args.negative_multiplier,
            rng=rng,
        )

    yaml_path = write_dataset_yaml(dataset_root)
    summary_path = output_root / "dataset_summary.json"
    summary_path.write_text(
        json.dumps(
            {
                "steel_sources": len(steel_paths),
                "negative_pool": len(negative_paths),
                "negative_source": negative_source,
                "splits": summary,
                "dataset_yaml": str(yaml_path),
            },
            indent=2,
        ),
        encoding="utf-8",
    )

    print(f"Dataset YAML written to: {yaml_path}")
    print(f"Summary written to: {summary_path}")


if __name__ == "__main__":
    main()