222 lines
8.2 KiB
Python
222 lines
8.2 KiB
Python
"""
|
|
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()
|