Refine dataset mask upload workflow

This commit is contained in:
2026-07-01 11:07:32 +08:00
parent 1e8339e237
commit afd6bd39ec
4 changed files with 206 additions and 102 deletions

View File

@@ -1,7 +1,12 @@
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, validate_dataset
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):
@@ -40,3 +45,33 @@ def test_validate_dataset_and_generate_yolo_yaml(tmp_path, monkeypatch):
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