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