"""Prepare mock test data from a folder of dog-identity subfolders. Input layout: input_dir/ dog_a/ img1.jpg img2.jpg ... dog_b/ img1.jpg ... Per dog folder, images are split into registration images (-> KnownDog) and holdout images (-> UnknownDog + found case). Writes known_dogs.csv, found_dogs.csv, and pairs.csv (eval only) to the output dir, with procedurally-generated owner/finder info. The found ZIP is set equal to the dog's known ZIP so default-radius matching works. Usage: python -m scripts.prepare_test_data --input-dir PATH [--output-dir PATH] \ [--holdout N] [--seed N] """ from __future__ import annotations import argparse import csv import random from pathlib import Path IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"} # Sample ZIPs that exist in data/zip_centroids.csv so radius matching is meaningful. SAMPLE_ZIPS = ["10001", "20001", "60601", "94103", "98101", "30301", "78701", "85001"] COLORS = ["black", "brown", "white", "tan", "golden", "brindle", "gray", "spotted"] SIZES = ["small", "medium", "large"] DOG_NAMES = ["Rex", "Bella", "Max", "Luna", "Charlie", "Daisy", "Cooper", "Lucy", "Buddy", "Sadie"] FIRST_NAMES = ["Alex", "Sam", "Jordan", "Taylor", "Casey", "Morgan", "Riley", "Jamie", "Pat", "Quinn"] LAST_NAMES = ["Smith", "Johnson", "Lee", "Garcia", "Brown", "Davis", "Martinez", "Clark", "Lewis", "Walker"] def _images_in(folder: Path) -> list[Path]: return sorted(p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS) def _holdout_count(n_images: int, requested: int | None) -> int: """Default 1 holdout; 2 if 5+ images and not specified. Always leave >=1 registration image.""" if requested is not None: h = requested else: h = 2 if n_images >= 5 else 1 return max(0, min(h, n_images - 1)) def _fake_person(rng: random.Random) -> tuple[str, str, str]: name = f"{rng.choice(FIRST_NAMES)} {rng.choice(LAST_NAMES)}" handle = name.lower().replace(" ", ".") email = f"{handle}{rng.randint(1, 999)}@example.com" phone = f"{rng.randint(200, 989)}-{rng.randint(200, 989)}-{rng.randint(1000, 9999)}" return name, email, phone def prepare(input_dir: Path, output_dir: Path, holdout: int | None, seed: int) -> dict: rng = random.Random(seed) subfolders = sorted(p for p in input_dir.iterdir() if p.is_dir()) known_rows: list[dict] = [] found_rows: list[dict] = [] pair_rows: list[dict] = [] skipped: list[str] = [] for folder in subfolders: images = _images_in(folder) if not images: skipped.append(folder.name) continue # Per-dog attributes (deterministic given seed + folder order). dog_name = rng.choice(DOG_NAMES) color = rng.choice(COLORS) size = rng.choice(SIZES) zip_code = rng.choice(SAMPLE_ZIPS) owner_name, owner_email, owner_phone = _fake_person(rng) finder_name, finder_email, finder_phone = _fake_person(rng) description = f"{size} {color} dog" n_hold = _holdout_count(len(images), holdout) shuffled = images[:] rng.shuffle(shuffled) holdouts = shuffled[:n_hold] registrations = shuffled[n_hold:] for img in registrations: known_rows.append( { "folder": folder.name, "image_file": img.name, "dog_name": dog_name, "color": color, "size": size, "description": description, "zip": zip_code, "owner_name": owner_name, "owner_email": owner_email, "owner_phone": owner_phone, } ) for img in holdouts: found_rows.append( { "folder": folder.name, "image_file": img.name, "description": f"found {description}", "color": color, "size": size, "found_zip": zip_code, # == known zip so default-radius matching works "current_location": "Local shelter", "finder_name": finder_name, "finder_email": finder_email, "finder_phone": finder_phone, } ) pair_rows.append( { "dog_folder": folder.name, "known_images": ";".join(i.name for i in registrations), "found_images": ";".join(i.name for i in holdouts), } ) output_dir.mkdir(parents=True, exist_ok=True) _write_csv(output_dir / "known_dogs.csv", known_rows, ["folder", "image_file", "dog_name", "color", "size", "description", "zip", "owner_name", "owner_email", "owner_phone"]) _write_csv(output_dir / "found_dogs.csv", found_rows, ["folder", "image_file", "description", "color", "size", "found_zip", "current_location", "finder_name", "finder_email", "finder_phone"]) _write_csv(output_dir / "pairs.csv", pair_rows, ["dog_folder", "known_images", "found_images"]) return { "dogs": len(pair_rows), "known_images": len(known_rows), "found_images": len(found_rows), "skipped_folders": skipped, } def _write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None: with path.open("w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) def main() -> None: parser = argparse.ArgumentParser(description="Prepare mock test data from dog-identity folders.") parser.add_argument("--input-dir", required=True, type=Path) parser.add_argument("--output-dir", type=Path, default=None) parser.add_argument("--holdout", type=int, default=None, help="Holdout images per dog (default 1; 2 if 5+ images)") parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() output_dir = args.output_dir or args.input_dir result = prepare(args.input_dir, output_dir, args.holdout, args.seed) print(f"Prepared {result['dogs']} dog(s): " f"{result['known_images']} registration + {result['found_images']} holdout image(s).") if result["skipped_folders"]: print(f"Skipped (no images): {', '.join(result['skipped_folders'])}") print(f"Wrote known_dogs.csv, found_dogs.csv, pairs.csv to {output_dir}") if __name__ == "__main__": main()