File size: 3,906 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 | """Tests for scripts/prepare_test_data.py — split ratios, reproducibility, CSV format, pairs."""
import csv
from pathlib import Path
from scripts.make_sample_images import make_image
from scripts.prepare_test_data import prepare
def _make_dataset_dir(tmp: Path, layout: dict[str, int]) -> Path:
"""layout: {folder_name: image_count}. Writes real JPEGs."""
root = tmp / "input"
for folder, count in layout.items():
d = root / folder
d.mkdir(parents=True)
for i in range(count):
(d / f"img{i}.jpg").write_bytes(make_image(hash((folder, i)) % 1000))
return root
def _read(path: Path) -> list[dict]:
with path.open(newline="", encoding="utf-8") as fh:
return list(csv.DictReader(fh))
def test_default_holdout_split(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 3, "dogB": 6})
out = tmp_path / "out"
prepare(root, out, holdout=None, seed=1)
known = _read(out / "known_dogs.csv")
found = _read(out / "found_dogs.csv")
by_folder = lambda rows, f: [r for r in rows if r["folder"] == f] # noqa: E731
# dogA has <5 images -> 1 holdout; dogB has 5+ -> 2 holdout.
assert len(by_folder(found, "dogA")) == 1
assert len(by_folder(known, "dogA")) == 2
assert len(by_folder(found, "dogB")) == 2
assert len(by_folder(known, "dogB")) == 4
def test_explicit_holdout(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogB": 6})
out = tmp_path / "out"
prepare(root, out, holdout=1, seed=1)
assert len(_read(out / "found_dogs.csv")) == 1
assert len(_read(out / "known_dogs.csv")) == 5
def test_reproducible_with_same_seed(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 4, "dogB": 5})
out1, out2 = tmp_path / "o1", tmp_path / "o2"
prepare(root, out1, holdout=None, seed=7)
prepare(root, out2, holdout=None, seed=7)
for fname in ("known_dogs.csv", "found_dogs.csv", "pairs.csv"):
assert (out1 / fname).read_text() == (out2 / fname).read_text()
def test_csv_columns_match_loader_contract(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 2})
out = tmp_path / "out"
prepare(root, out, holdout=1, seed=1)
known_cols = _read(out / "known_dogs.csv")[0].keys()
found_cols = _read(out / "found_dogs.csv")[0].keys()
assert set(known_cols) == {
"folder", "image_file", "dog_name", "color", "size", "description",
"zip", "owner_name", "owner_email", "owner_phone",
}
assert set(found_cols) == {
"folder", "image_file", "description", "color", "size", "found_zip",
"current_location", "finder_name", "finder_email", "finder_phone",
}
def test_found_zip_equals_known_zip(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 3})
out = tmp_path / "out"
prepare(root, out, holdout=1, seed=3)
known_zip = _read(out / "known_dogs.csv")[0]["zip"]
found_zip = _read(out / "found_dogs.csv")[0]["found_zip"]
assert known_zip == found_zip # so default-radius matching works
def test_pairs_link_known_and_found(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 4})
out = tmp_path / "out"
prepare(root, out, holdout=1, seed=5)
pairs = _read(out / "pairs.csv")
assert len(pairs) == 1
row = pairs[0]
known_imgs = set(row["known_images"].split(";"))
found_imgs = set(row["found_images"].split(";"))
assert len(found_imgs) == 1
assert len(known_imgs) == 3
assert known_imgs.isdisjoint(found_imgs) # an image is either registration or holdout
assert known_imgs | found_imgs == {"img0.jpg", "img1.jpg", "img2.jpg", "img3.jpg"}
def test_skips_empty_folder(tmp_path):
root = _make_dataset_dir(tmp_path, {"dogA": 2})
(root / "empty_dog").mkdir()
out = tmp_path / "out"
result = prepare(root, out, holdout=1, seed=1)
assert "empty_dog" in result["skipped_folders"]
assert result["dogs"] == 1
|