Files
ai-training/tests/test_annotation_classes.py
T

80 lines
2.2 KiB
Python

import re
import sys
import types
import pytest
from dto.annotationClass import AnnotationClass
def _stub_train_imports():
if getattr(_stub_train_imports, "_done", False):
return
for _name in ("ultralytics", "boto3", "netron", "requests"):
if _name not in sys.modules:
sys.modules[_name] = types.ModuleType(_name)
sys.modules["ultralytics"].YOLO = type("YOLO", (), {})
sys.modules["boto3"].client = lambda *a, **k: None
_stub_train_imports._done = True
def _name_lines_under_names(text):
lines = text.splitlines()
out = []
in_block = False
for line in lines:
s = line.strip()
if s == "names:":
in_block = True
continue
if s.startswith("nc:"):
break
if in_block and s.startswith("-"):
out.append(s)
return out
_PLACEHOLDER_RE = re.compile(r"^-\s+Class-\d+\s*$")
@pytest.fixture
def data_yaml_text(monkeypatch, tmp_path, fixture_classes_json):
_stub_train_imports()
import train
monkeypatch.setattr(train, "today_dataset", str(tmp_path))
train.create_yaml()
return (tmp_path / "data.yaml").read_text(encoding="utf-8")
def test_bt_cls_01_base_classes(fixture_classes_json):
d = AnnotationClass.read_json()
norm = {k: d[k] for k in range(17)}
assert len(norm) == 17
assert len({v.id for v in norm.values()}) == 17
def test_bt_cls_02_weather_expansion(fixture_classes_json):
d = AnnotationClass.read_json()
assert d[0].name == "ArmorVehicle"
assert d[20].name == "ArmorVehicle(Wint)"
assert d[40].name == "ArmorVehicle(Night)"
@pytest.mark.resource_limit
def test_bt_cls_03_yaml_generation(data_yaml_text):
text = data_yaml_text
assert "nc: 80" in text
names = _name_lines_under_names(text)
placeholders = [ln for ln in names if _PLACEHOLDER_RE.match(ln)]
named = [ln for ln in names if not _PLACEHOLDER_RE.match(ln)]
assert len(names) == 80
assert len(placeholders) == 29
assert len(named) == 51
@pytest.mark.resource_limit
def test_rl_cls_01_total_class_count(data_yaml_text):
names = _name_lines_under_names(data_yaml_text)
assert len(names) == 80