Shaun ZhouandClaude Opus 5
36fff984b7
原先按 split 和结果类型（ok/missed/extra/none）分目录，层级太深。改成 <类别>/<图名>.png，与 runs/boxes_shap 的结构一致。类别取真值里框最多的那个； 负样本没有真值框，按文件名前缀归类（<类别>_<序号> 是 prepare_label_studio 的命名）， 认不出来放 background/。 split 与结果类型仍记在 index.csv 的列里，排序筛选即可。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
"""可视化脚本里不依赖模型的部分。"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("yaml")
|
|
import yaml # noqa: E402
|
|
|
|
sys.path.insert(0, "scripts")
|
|
import visualize_detections as vd # noqa: E402
|
|
|
|
from engine.detection import Box, from_yolo, match_boxes # noqa: E402
|
|
|
|
|
|
def _match(predictions, truth):
|
|
return match_boxes(predictions, truth)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"truth,predictions,expected",
|
|
[
|
|
([], [], "none"), # 负样本判对
|
|
([from_yolo(0, .5, .5, .2, .2)], [], "missed"), # 完全没检出
|
|
([], [Box(0, .4, .4, .6, .6, .9)], "extra"), # 背景上凭空框出
|
|
([from_yolo(0, .5, .5, .2, .2)], [Box(0, .41, .41, .59, .59, .9)], "ok"),
|
|
],
|
|
)
|
|
def test_outcome_categories(truth, predictions, expected):
|
|
assert vd.outcome(len(truth), len(predictions), _match(predictions, truth)) == expected
|
|
|
|
|
|
def test_missed_takes_priority_over_extra():
|
|
"""一张图同时有漏检和误检时归为 missed——漏检更该先看。"""
|
|
truth = [from_yolo(0, .2, .2, .1, .1)]
|
|
predictions = [Box(0, .7, .7, .9, .9, .9)]
|
|
assert vd.outcome(1, 1, _match(predictions, truth)) == "missed"
|
|
|
|
|
|
def test_label_path_swaps_images_for_labels():
|
|
image = Path("/data/film/images/val/blackdot_0001.png")
|
|
assert vd.label_path_for(image) == Path("/data/film/labels/val/blackdot_0001.txt")
|
|
|
|
|
|
def test_label_path_uses_the_last_images_segment():
|
|
"""路径里多处出现 images 时，只换最靠后的那一段。"""
|
|
image = Path("/srv/images/film/images/train/a.png")
|
|
assert vd.label_path_for(image) == Path("/srv/images/film/labels/train/a.txt")
|
|
|
|
|
|
@pytest.fixture
|
|
def dataset(tmp_path):
|
|
for split in ("train", "val"):
|
|
(tmp_path / "images" / split).mkdir(parents=True)
|
|
(tmp_path / "images" / split / f"{split}_0001.png").write_bytes(b"")
|
|
data_yaml = tmp_path / "data.yaml"
|
|
data_yaml.write_text(
|
|
yaml.safe_dump({
|
|
"path": str(tmp_path), "train": "images/train", "val": "images/val",
|
|
"nc": 2, "names": {0: "blackdot", 1: "grease"},
|
|
}),
|
|
encoding="utf-8",
|
|
)
|
|
return data_yaml
|
|
|
|
|
|
def test_load_dataset(dataset):
|
|
classes, splits = vd.load_dataset(dataset)
|
|
assert classes == ["blackdot", "grease"]
|
|
assert sorted(splits) == ["train", "val"]
|
|
assert splits["val"][0].name == "val_0001.png"
|
|
|
|
|
|
def test_load_dataset_skips_a_split_pointing_at_the_same_directory(dataset):
|
|
"""未划分的数据集 train 与 val 指向同一目录，不该把同一批图跑两遍。"""
|
|
meta = yaml.safe_load(dataset.read_text(encoding="utf-8"))
|
|
meta["val"] = meta["train"]
|
|
dataset.write_text(yaml.safe_dump(meta), encoding="utf-8")
|
|
_, splits = vd.load_dataset(dataset)
|
|
assert list(splits) == ["train"]
|
|
|
|
|
|
def test_load_dataset_reports_a_missing_directory(dataset):
|
|
meta = yaml.safe_load(dataset.read_text(encoding="utf-8"))
|
|
meta["val"] = "images/nope"
|
|
dataset.write_text(yaml.safe_dump(meta), encoding="utf-8")
|
|
with pytest.raises(FileNotFoundError, match="不存在"):
|
|
vd.load_dataset(dataset)
|
|
|
|
|
|
# ---------------- 产出的目录结构 ----------------
|
|
|
|
|
|
def test_class_dir_uses_the_dominant_ground_truth_class():
|
|
classes = ["blackdot", "grease", "normal"]
|
|
truth = [from_yolo(1, .5, .5, .2, .2), from_yolo(0, .1, .1, .1, .1), from_yolo(1, .8, .8, .1, .1)]
|
|
assert vd.class_dir_for(truth, classes, "grease_0007") == "grease"
|
|
|
|
|
|
def test_class_dir_falls_back_to_the_filename_prefix():
|
|
"""负样本没有真值框，按 <类别>_<序号> 的命名归类。"""
|
|
classes = ["blackdot", "grease", "normal"]
|
|
assert vd.class_dir_for([], classes, "normal_1234") == "normal"
|
|
|
|
|
|
def test_class_dir_falls_back_to_background_for_unknown_names():
|
|
classes = ["blackdot", "grease", "normal"]
|
|
assert vd.class_dir_for([], classes, "IMG_20260101") == "background"
|
|
assert vd.class_dir_for([], classes, "nofix") == "background"
|