diff --git a/backend/app/modules/dataset/service.py b/backend/app/modules/dataset/service.py index accd6ec..a05b26a 100644 --- a/backend/app/modules/dataset/service.py +++ b/backend/app/modules/dataset/service.py @@ -22,6 +22,7 @@ ARCHIVE_ALIASES = { "masks": {"mask", "masks", "label", "labels", "gt", "annotation", "annotations"}, "labels": {"label", "labels", "txt", "annotation", "annotations"}, } +MASK_STEM_SUFFIXES = ("-mask", "_mask", "-label", "_label", "-gt", "_gt", "-seg", "_seg") def uploads_root() -> Path: @@ -44,6 +45,16 @@ def safe_filename(value: str | None) -> str: return stem +def upload_exts_for_kind(kind: str) -> set[str]: + if kind == "images": + return IMAGE_EXTS + if kind == "masks": + return IMAGE_EXTS | LABEL_EXTS + if kind == "labels": + return LABEL_EXTS + return set() + + def is_archive_name(value: str | None) -> bool: name = (value or "").lower() return any(name.endswith(suffix) for suffix in ARCHIVE_SUFFIXES) @@ -88,9 +99,19 @@ def safe_archive_member(raw_name: str, kind: str) -> Path | None: return None safe_parts = [slugify(part) for part in raw_parts[:-1]] filename = safe_filename(raw_parts[-1]) + if Path(filename).suffix.lower() not in upload_exts_for_kind(kind): + return None return Path(*safe_parts, filename) if safe_parts else Path(filename) +def normalize_mask_stem(stem: str) -> str: + lower = stem.lower() + for suffix in MASK_STEM_SUFFIXES: + if lower.endswith(suffix) and len(stem) > len(suffix): + return stem[: -len(suffix)] + return stem + + def save_archive_member(target: Path, member_name: str, kind: str, source) -> dict | None: relative = safe_archive_member(member_name, kind) if relative is None: @@ -147,12 +168,13 @@ def _iter_files(root: Path) -> Iterable[Path]: return sorted(path for path in root.rglob("*") if path.is_file()) -def _stem_map(paths: Iterable[Path], exts: set[str] | None = None) -> dict[str, Path]: +def _stem_map(paths: Iterable[Path], exts: set[str] | None = None, normalize=None) -> dict[str, Path]: result: dict[str, Path] = {} for path in paths: if exts and path.suffix.lower() not in exts: continue - result.setdefault(path.stem, path) + stem = normalize(path.stem) if normalize else path.stem + result.setdefault(stem, path) return result @@ -234,7 +256,7 @@ def validate_dataset(name: str) -> dict: mask_files = list(_iter_files(root / "masks")) images = _stem_map(image_files, IMAGE_EXTS) labels = _stem_map(label_files, LABEL_EXTS) - masks = _stem_map(mask_files, IMAGE_EXTS) + masks = _stem_map(mask_files, IMAGE_EXTS, normalize_mask_stem) image_stems = set(images) label_stems = set(labels) diff --git a/backend/tests/test_dataset_service.py b/backend/tests/test_dataset_service.py index 9b43999..d968321 100644 --- a/backend/tests/test_dataset_service.py +++ b/backend/tests/test_dataset_service.py @@ -31,7 +31,7 @@ def test_validate_dataset_and_generate_yolo_yaml(tmp_path, monkeypatch): 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) + cv2.imwrite(str(root / "masks" / "sample-mask.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") @@ -57,6 +57,7 @@ def test_upload_zip_extracts_matching_dataset_kind(tmp_path, monkeypatch): image_zip = io.BytesIO() with zipfile.ZipFile(image_zip, "w") as archive: archive.writestr("bundle/images/sample.png", b"image") + archive.writestr("bundle/images/manifest.csv", b"name\nsample.png\n") archive.writestr("bundle/masks/sample.png", b"mask") image_zip.seek(0) image_upload = UploadFile(filename="images.zip", file=image_zip) @@ -65,13 +66,13 @@ def test_upload_zip_extracts_matching_dataset_kind(tmp_path, monkeypatch): 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") + archive.writestr("bundle/masks/sample-mask.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"] + assert [item["relative_path"] for item in mask_result["saved"]] == ["var/uploads/datasets/case_zip/masks/sample-mask.png"] described = describe_dataset("case_zip") assert described["counts"]["images"] == 1 assert described["counts"]["masks"] == 1