Files
AareLC_train/tests/test_import_studio_export.py
T

67 lines
3.5 KiB
Python

"""tools/import_studio_export.py: images once in the pool, symlinked per dataset."""
import json
from pathlib import Path
import pytest
from tools import import_studio_export as imp
YAML = "path: /x\ntrain: train/images\nval: val/images\n\nnames:\n 0: loop_all\n 1: pin\n"
def make_export(root: Path, name: str, train=("tpi_1", "tpi_2"), val=("tpi_3",), missing_label=()):
exp = root / name
for split, keys in (("train", train), ("val", val)):
(exp / split / "images").mkdir(parents=True)
(exp / split / "labels").mkdir(parents=True)
for k in keys:
(exp / split / "images" / f"{k}.jpg").write_bytes(b"\xff\xd8" + k.encode() * 20)
if k not in missing_label:
(exp / split / "labels" / f"{k}.txt").write_text("0 0.1 0.1 0.2 0.1 0.2 0.2\n" if k != "tpi_2" else "")
(exp / "dataset.yaml").write_text(YAML)
(exp / "manifest.csv").write_text("key,tpid,split\n" + "\n".join(f"{k},{k[4:]},train" for k in train) + "\n")
return exp
def test_import_builds_pool_and_symlinks(tmp_path):
repo = tmp_path / "repo"; repo.mkdir()
exp = make_export(tmp_path, "export_a", missing_label=("tpi_3",))
src = imp.import_export(exp, "2026-10-02_alc", root=repo)
ds = repo / "data" / "datasets" / "2026-10-02_alc"
pool = repo / "data" / "pool" / "images"
assert sorted(p.name for p in pool.iterdir()) == ["tpi_1.jpg", "tpi_2.jpg", "tpi_3.jpg"]
link = ds / "images" / "train" / "tpi_1.jpg"
assert link.is_symlink() and link.resolve() == (pool / "tpi_1.jpg").resolve()
assert (ds / "labels" / "train" / "tpi_1.txt").read_text().startswith("0 0.1")
assert (ds / "labels" / "val" / "tpi_3.txt").read_text() == "" # missing label -> background
assert src["counts"]["train"] == {"images": 2, "labels": 2, "empty": 1}
assert src["counts"]["val"] == {"images": 1, "labels": 0, "empty": 1}
assert src["new_pool_images"] == 3 and src["classes"] == {0: "loop_all", 1: "pin"}
assert "path: ." in (ds / "dataset.yaml").read_text() and (ds / "manifest.csv").exists()
assert json.loads((ds / "source.json").read_text())["name"] == "2026-10-02_alc"
def test_second_dataset_reuses_the_pool_and_names_are_immutable(tmp_path):
repo = tmp_path / "repo"; repo.mkdir()
imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo)
src = imp.import_export(make_export(tmp_path, "b", train=("tpi_1", "tpi_2", "tpi_9"), val=("tpi_3",)),
"2026-10-09_alc", root=repo)
assert src["new_pool_images"] == 1
assert len(list((repo / "data" / "pool" / "images").iterdir())) == 4
with pytest.raises(SystemExit):
imp.import_export(make_export(tmp_path, "c"), "2026-10-02_alc", root=repo)
def test_gold_import(tmp_path):
repo = tmp_path / "repo"; repo.mkdir()
gold = tmp_path / "gold"
(gold / "images").mkdir(parents=True); (gold / "labels").mkdir()
(gold / "images" / "tpi_7.jpg").write_bytes(b"\xff\xd8gold")
(gold / "labels" / "tpi_7.txt").write_text("1 0.5 0.5 0.6 0.5 0.6 0.6\n")
(gold / "gold_manifest.json").write_text(json.dumps({"keys": ["tpi_7"]}))
src = imp.import_export(make_export(tmp_path, "a"), "2026-10-02_alc", root=repo, gold_dir=gold)
assert src["gold"]["images"] == 1 and src["gold"]["manifest_sha256"]
assert (repo / "data" / "gold" / "images" / "tpi_7.jpg").is_symlink()
assert (repo / "data" / "gold" / "labels" / "tpi_7.txt").exists()
assert (repo / "data" / "pool" / "images" / "tpi_7.jpg").exists()