Add dataset QA and custom YOLO training flow
This commit is contained in:
@@ -7,9 +7,9 @@ def test_catalog_contains_required_capabilities():
|
||||
for task in [
|
||||
"dataset.video_frames",
|
||||
"segmodel.train",
|
||||
"yolo.train_custom",
|
||||
"yolo.predict",
|
||||
"mmseg.flops_fps",
|
||||
"analysis.all",
|
||||
]:
|
||||
assert task in TASK_TYPES
|
||||
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
from app.modules.dataset.service import create_dataset, describe_dataset
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from app.modules.dataset.service import create_dataset, describe_dataset, generate_yolo_dataset_yaml, validate_dataset
|
||||
|
||||
|
||||
def test_create_dataset_layout(tmp_path, monkeypatch):
|
||||
@@ -11,3 +14,29 @@ def test_create_dataset_layout(tmp_path, monkeypatch):
|
||||
assert created["counts"] == {"images": 0, "labels": 0, "masks": 0}
|
||||
described = describe_dataset("case_01")
|
||||
assert described["layout"]["images"].endswith("images")
|
||||
|
||||
|
||||
def test_validate_dataset_and_generate_yolo_yaml(tmp_path, monkeypatch):
|
||||
from types import SimpleNamespace
|
||||
from app.modules.dataset import service
|
||||
|
||||
monkeypatch.setattr(service, "settings", SimpleNamespace(project_root=tmp_path))
|
||||
create_dataset("case yolo", "demo")
|
||||
root = service.dataset_dir("case_yolo")
|
||||
image = np.zeros((32, 32, 3), dtype=np.uint8)
|
||||
mask = np.zeros((32, 32), dtype=np.uint8)
|
||||
cv2.imwrite(str(root / "images" / "sample.png"), image)
|
||||
cv2.imwrite(str(root / "masks" / "sample.png"), mask)
|
||||
(root / "labels" / "sample.txt").write_text("0 0.25 0.25 0.75 0.25 0.75 0.75 0.25 0.75\n", encoding="utf-8")
|
||||
|
||||
validation = validate_dataset("case_yolo")
|
||||
|
||||
assert validation["ready"] == {"yolo": True, "mask": True, "any": True}
|
||||
assert validation["pairs"]["image_label"] == 1
|
||||
assert validation["classes"] == [0]
|
||||
|
||||
generated = generate_yolo_dataset_yaml("case_yolo", ["tool"])
|
||||
|
||||
assert generated["relative_path"] == "var/uploads/datasets/case_yolo/dataset.yaml"
|
||||
assert "nc: 1" in generated["content"]
|
||||
assert "0: tool" in generated["content"]
|
||||
|
||||
Reference in New Issue
Block a user