File size: 6,661 Bytes
de1e3fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """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()
|