File size: 1,742 Bytes
f8c73f9 | 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 | from pathlib import Path
import pytest
from adam.cnn_reviewer import _labeled_paths, reviewer_checkpoint, score_frames, train_reviewer
def test_reviewer_checkpoint_is_adam_owned_and_dataset_specific(tmp_path: Path) -> None:
first = reviewer_checkpoint(tmp_path, tmp_path / "dataset one")
second = reviewer_checkpoint(tmp_path, tmp_path / "dataset two")
assert first.parent == tmp_path / "data" / "cnn_reviewers"
assert first != second
assert first.suffix == ".pt"
def test_reviewer_uses_only_readable_image_keep_reject_labels(tmp_path: Path) -> None:
kept = tmp_path / "keep.png"; kept.write_bytes(b"not decoded here")
rejected = tmp_path / "reject.jpg"; rejected.write_bytes(b"not decoded here")
missing = tmp_path / "missing.png"
paths, labels = _labeled_paths({str(kept): "keep", str(rejected): "reject", str(missing): "keep", "x": "unreviewed"})
assert paths == [kept, rejected]
assert labels == [1.0, 0.0]
def test_reviewer_trains_and_scores_without_touching_decisions(tmp_path: Path) -> None:
pytest.importorskip("torch")
from PIL import Image
decisions = {}
for index in range(8):
path = tmp_path / f"keep_{index}.png"
Image.new("RGB", (32, 32), (20, 220, 30)).save(path)
decisions[str(path)] = "keep"
for index in range(8):
path = tmp_path / f"reject_{index}.png"
Image.new("RGB", (32, 32), (220, 20, 30)).save(path)
decisions[str(path)] = "reject"
result = train_reviewer(tmp_path, tmp_path, decisions, epochs=1)
scores = score_frames(result.checkpoint, tmp_path)
assert Path(result.checkpoint).is_file()
assert len(scores) == 16
assert decisions[str(tmp_path / "keep_0.png")] == "keep"
|