File size: 4,159 Bytes
ebeed67 | 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 | from __future__ import annotations
import io
import tarfile
import pytest
from camstack_train.dataset import (
coco_80, coco_index_for, extract_archive, frames_from_manifest, verify_archive, write_layout,
)
from camstack_train.errors import EXIT_DATASET, TrainerError
from .conftest import frame, subject
def test_coco_mapping_is_honest():
names = coco_80()
assert coco_index_for("person", None, names) == 0
assert coco_index_for("vehicle", "Car", names) == names.index("car")
assert coco_index_for("animal", "dog", names) == names.index("dog")
assert coco_index_for("vehicle", None, names) is None # no class to give it
assert coco_index_for("animal", "car", names) is None # label of another macro
def test_boxes_become_yolo_lines_and_model_errors_are_not_positives(make_export, tmp_path):
frames = [
frame("f1", 615, 0, [subject("person"), subject("vehicle", "car", x=0.5), subject("person", kind="model_error", x=0.6)]),
frame("f2", 615, 0, []), # complete, empty: a hard negative
frame("f3", 617, 0, [subject("animal")]), # animal without a COCO label: unmappable
frame("f4", 617, 0, [subject("package")], complete=False), # incomplete, no detector subject…
frame("f5", 617, 0, [], complete=False), # incomplete, nothing: never trains
]
path, digest = make_export(frames)
verify_archive(path, digest)
manifest = extract_archive(path, tmp_path / "x")
parsed = {f.frame_id: f for f in frames_from_manifest(manifest)}
f1 = parsed["f1"]
assert [b.cls_index for b in f1.labels] == [0, coco_80().index("car")]
assert f1.model_errors == 1 and f1.trainable
assert parsed["f2"].trainable and parsed["f2"].labels == ()
assert not parsed["f3"].trainable and parsed["f3"].unmappable == 1
assert not parsed["f4"].trainable # a package box is not a root-detector subject
assert not parsed["f5"].trainable
box = f1.labels[0]
assert (round(box.cx, 3), round(box.cy, 3), round(box.w, 3), round(box.h, 3)) == (0.25, 0.4, 0.3, 0.4)
counts = write_layout({"train": [f for f in parsed.values() if f.trainable]}, tmp_path / "x", tmp_path / "yolo")
assert counts.frames["train"] == 2
label_files = sorted((tmp_path / "yolo" / "labels" / "train").glob("*.txt"))
assert [p.name for p in label_files] == ["000001.txt", "000002.txt"] # neutral names
assert label_files[0].read_text().splitlines()[0].startswith("0 0.250000 0.400000")
assert label_files[1].read_text() == ""
def test_incomplete_frame_with_a_box_trains(make_export, tmp_path):
path, _ = make_export([frame("f1", 1, 0, [subject("person")], complete=False)])
f = frames_from_manifest(extract_archive(path, tmp_path / "x"))[0]
assert f.trainable and not f.complete
def test_out_of_frame_box_is_clamped_or_dropped():
names = coco_80()
from camstack_train.dataset import frame_from_entry
entry = frame("f", 1, 0, [
{**subject("person"), "bbox": {"x": -0.2, "y": 0.0, "w": 0.4, "h": 0.5}},
{**subject("person"), "bbox": {"x": 1.2, "y": 0.0, "w": 0.4, "h": 0.5}},
])
f = frame_from_entry(entry, names)
assert len(f.labels) == 1 and round(f.labels[0].w, 3) == 0.2
def test_sha256_mismatch_is_a_dataset_error(make_export):
path, _ = make_export([frame("f1", 1, 0, [])])
with pytest.raises(TrainerError) as err:
verify_archive(path, "1" * 64)
assert err.value.exit_code == EXIT_DATASET and err.value.code == "dataset-sha256"
def test_unknown_export_schema_refused(make_export, tmp_path):
path, _ = make_export([frame("f1", 1, 0, [])], schema="camstack.retrain-export/v9")
with pytest.raises(TrainerError) as err:
extract_archive(path, tmp_path / "x")
assert err.value.code == "dataset-schema"
def test_extraction_skips_unsafe_members(tmp_path, make_export):
path, _ = make_export([frame("f1", 1, 0, [])])
with tarfile.open(path, "a") as tar:
info = tarfile.TarInfo("../evil.jpg")
info.size = 1
tar.addfile(info, io.BytesIO(b"x"))
extract_archive(path, tmp_path / "x")
assert not (tmp_path / "evil.jpg").exists()
|