""" Script 1b — Dataset Splitter ============================== Takes your labelled data/ folder (one subfolder per class) and splits it into the train / val / test structure expected by the training script. The split is STRATIFIED — every class gets the same train/val/test ratio, so rare classes are not accidentally wiped out of val or test. Input layout (after you've labelled crops from script 1) --------------------------------------------------------- data/ cola_can/ ← labelled class folders img001_crop0000.jpg img002_crop0003.jpg ... pepsi_can/ ... unknown/ ← IGNORED automatically (unlabelled crops) ... Output layout ------------- crops_dataset/ train/ cola_can/ ... pepsi_can/ ... val/ cola_can/ ... pepsi_can/ ... test/ (only if --test_split > 0) cola_can/ ... pepsi_can/ ... Files are COPIED by default (--move to use move instead). Original data/ folder is left intact unless --move is used. Usage ----- python 1b_split_dataset.py \ --data_dir data \ --output_dir crops_dataset \ --val_split 0.15 \ --test_split 0.05 \ --seed 42 Then run: python 2_train_classifier.py --data_dir crops_dataset """ import argparse import json import random import shutil from collections import defaultdict from pathlib import Path # ─────────────────────────── args ──────────────────────────────────────────── def parse_args(): p = argparse.ArgumentParser(description="Stratified train/val/test splitter") p.add_argument("--data_dir", required=True, help="Labelled data root (subfolders = classes)") p.add_argument("--output_dir", default="crops_dataset", help="Where to write the split dataset") p.add_argument("--val_split", type=float, default=0.15, help="Fraction for validation (default 0.15)") p.add_argument("--test_split", type=float, default=0.05, help="Fraction for test — 0 to skip") p.add_argument("--seed", type=int, default=42) p.add_argument("--move", action="store_true", help="Move files instead of copying (saves disk space)") p.add_argument("--min_per_class", type=int, default=2, help="Warn if a class has fewer than this many images after splitting") return p.parse_args() # ─────────────────────────── helpers ───────────────────────────────────────── IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"} IGNORE_DIRS = {"unknown", "train", "val", "test"} # skip these silently def collect_classes(data_dir: Path) -> dict[str, list[Path]]: """ Scan *data_dir* for class subfolders and return {class_name: [image_paths]}. Folders named 'unknown', 'train', 'val', or 'test' are silently skipped. """ classes: dict[str, list[Path]] = {} for folder in sorted(data_dir.iterdir()): if not folder.is_dir(): continue if folder.name.lower() in IGNORE_DIRS: print(f" ⚙ Skipping '{folder.name}/' (reserved name)") continue images = sorted(f for f in folder.iterdir() if f.suffix.lower() in IMAGE_EXTS) if images: classes[folder.name] = images else: print(f" ⚠ '{folder.name}/' has no images — skipped") return classes def stratified_split( images: list[Path], val_frac: float, test_frac: float, rng: random.Random, ) -> tuple[list[Path], list[Path], list[Path]]: """ Split *images* into (train, val, test) respecting the requested fractions. Guarantees at least 1 image in val and test if the class is large enough. """ imgs = images[:] rng.shuffle(imgs) n = len(imgs) n_test = max(1, round(n * test_frac)) if test_frac > 0 else 0 n_val = max(1, round(n * val_frac)) if val_frac > 0 else 0 # Ensure we don't eat into train entirely for tiny classes n_test = min(n_test, max(0, n - 2)) n_val = min(n_val, max(0, n - n_test - 1)) test = imgs[:n_test] val = imgs[n_test : n_test + n_val] train = imgs[n_test + n_val:] return train, val, test def transfer(src: Path, dst: Path, move: bool) -> None: dst.parent.mkdir(parents=True, exist_ok=True) if move: shutil.move(str(src), str(dst)) else: shutil.copy2(str(src), str(dst)) # ─────────────────────────── main ──────────────────────────────────────────── def main(): args = parse_args() rng = random.Random(args.seed) data_dir = Path(args.data_dir) out_dir = Path(args.output_dir) if not data_dir.exists(): raise FileNotFoundError(f"data_dir not found: {data_dir}") # ── Collect classes ─────────────────────────────────────────────────────── print(f"\nScanning: {data_dir.resolve()}") classes = collect_classes(data_dir) if not classes: raise RuntimeError( f"No labelled class folders found in {data_dir}.\n" "Create subfolders named after your product classes and move crops into them." ) num_classes = len(classes) total_imgs = sum(len(v) for v in classes.values()) print(f"\nFound {num_classes} classes | {total_imgs:,} total images\n") # ── Split each class separately (stratified) ────────────────────────────── split_counts: dict[str, dict[str, int]] = defaultdict(lambda: {"train": 0, "val": 0, "test": 0}) warnings = [] for cls_name, images in classes.items(): train, val, test = stratified_split(images, args.val_split, args.test_split, rng) for split_name, split_imgs in [("train", train), ("val", val), ("test", test)]: if not split_imgs: continue for src in split_imgs: dst = out_dir / split_name / cls_name / src.name transfer(src, dst, move=args.move) split_counts[cls_name][split_name] = len(split_imgs) if len(train) < args.min_per_class: warnings.append( f" ⚠ '{cls_name}' has only {len(train)} training image(s) — " "consider collecting more data" ) # ── Write class manifest ────────────────────────────────────────────────── class_list = sorted(classes.keys()) class_map = {i: c for i, c in enumerate(class_list)} manifest_path = out_dir / "class_names.json" manifest_path.write_text(json.dumps(class_map, indent=2)) # ── Print summary table ─────────────────────────────────────────────────── col_w = max(len(c) for c in classes) + 2 header = f" {'Class':<{col_w}} {'Train':>7} {'Val':>7} {'Test':>7} {'Total':>7}" print("=" * len(header)) print(header) print("-" * len(header)) g_train = g_val = g_test = 0 for cls_name in sorted(split_counts): c = split_counts[cls_name] tot = c["train"] + c["val"] + c["test"] print(f" {cls_name:<{col_w}} {c['train']:>7} {c['val']:>7} {c['test']:>7} {tot:>7}") g_train += c["train"]; g_val += c["val"]; g_test += c["test"] print("-" * len(header)) print(f" {'TOTAL':<{col_w}} {g_train:>7} {g_val:>7} {g_test:>7} {g_train+g_val+g_test:>7}") print("=" * len(header)) if warnings: print() for w in warnings: print(w) action = "moved" if args.move else "copied" print(f"\n {action.capitalize()} to : {out_dir.resolve()}") print(f" class_names.json: {manifest_path}") print(f"\n Next step:") print(f" python 2_train_classifier.py --data_dir {out_dir}") if __name__ == "__main__": main()