78 lines
3.2 KiB
Python
78 lines
3.2 KiB
Python
import asyncio
|
|
import io
|
|
import zipfile
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from fastapi import UploadFile
|
|
|
|
from app.modules.dataset.service import create_dataset, describe_dataset, generate_yolo_dataset_yaml, save_upload, validate_dataset
|
|
|
|
|
|
def test_create_dataset_layout(tmp_path, monkeypatch):
|
|
from types import SimpleNamespace
|
|
from app.modules.dataset import service
|
|
|
|
monkeypatch.setattr(service, "settings", SimpleNamespace(project_root=tmp_path))
|
|
created = create_dataset("case 01", "demo")
|
|
assert created["name"] == "case_01"
|
|
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"]
|
|
|
|
|
|
def test_upload_zip_extracts_matching_dataset_kind(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 zip", "demo")
|
|
|
|
image_zip = io.BytesIO()
|
|
with zipfile.ZipFile(image_zip, "w") as archive:
|
|
archive.writestr("bundle/images/sample.png", b"image")
|
|
archive.writestr("bundle/masks/sample.png", b"mask")
|
|
image_zip.seek(0)
|
|
image_upload = UploadFile(filename="images.zip", file=image_zip)
|
|
image_result = asyncio.run(save_upload("case_zip", "images", [image_upload]))
|
|
|
|
mask_zip = io.BytesIO()
|
|
with zipfile.ZipFile(mask_zip, "w") as archive:
|
|
archive.writestr("bundle/images/sample.png", b"image")
|
|
archive.writestr("bundle/masks/sample.png", b"mask")
|
|
mask_zip.seek(0)
|
|
mask_upload = UploadFile(filename="masks.zip", file=mask_zip)
|
|
mask_result = asyncio.run(save_upload("case_zip", "masks", [mask_upload]))
|
|
|
|
assert [item["relative_path"] for item in image_result["saved"]] == ["var/uploads/datasets/case_zip/images/sample.png"]
|
|
assert [item["relative_path"] for item in mask_result["saved"]] == ["var/uploads/datasets/case_zip/masks/sample.png"]
|
|
described = describe_dataset("case_zip")
|
|
assert described["counts"]["images"] == 1
|
|
assert described["counts"]["masks"] == 1
|