Refine dataset mask upload workflow

backend/app/modules/dataset/service.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 import json
 import re
 import shutil
+import tarfile
+import zipfile
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Iterable
@@ -14,6 +16,12 @@ from ...config import settings
 DATASET_KINDS = ("images", "labels", "masks")
 IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
 LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}
+ARCHIVE_SUFFIXES = (".zip", ".tar", ".tar.gz", ".tgz")
+ARCHIVE_ALIASES = {
+    "images": {"image", "images", "img", "imgs", "ori", "original", "originals"},
+    "masks": {"mask", "masks", "label", "labels", "gt", "annotation", "annotations"},
+    "labels": {"label", "labels", "txt", "annotation", "annotations"},
+}
 
 
 def uploads_root() -> Path:
@@ -36,6 +44,68 @@ def safe_filename(value: str | None) -> str:
     return stem
 
 
+def is_archive_name(value: str | None) -> bool:
+    name = (value or "").lower()
+    return any(name.endswith(suffix) for suffix in ARCHIVE_SUFFIXES)
+
+
+def unique_path(path: Path) -> Path:
+    if not path.exists():
+        return path
+    stem = path.stem
+    suffix = path.suffix
+    counter = 1
+    while True:
+        candidate = path.with_name(f"{stem}_{counter}{suffix}")
+        if not candidate.exists():
+            return candidate
+        counter += 1
+
+
+def safe_archive_member(raw_name: str, kind: str) -> Path | None:
+    normalized = raw_name.replace("\\", "/")
+    if not normalized or normalized.endswith("/"):
+        return None
+    raw_parts = [part for part in normalized.split("/") if part not in {"", "."}]
+    if not raw_parts or any(part == ".." for part in raw_parts):
+        raise ValueError(f"unsafe archive member path: {raw_name}")
+    if Path(normalized).is_absolute():
+        raise ValueError(f"absolute archive member path is not allowed: {raw_name}")
+    lower_parts = [part.lower() for part in raw_parts]
+    if lower_parts[0] in {"__macosx", ".ds_store"}:
+        return None
+
+    target_aliases = ARCHIVE_ALIASES.get(kind, {kind})
+    other_aliases = set().union(*(aliases for item, aliases in ARCHIVE_ALIASES.items() if item != kind))
+    target_index = next((index for index, part in enumerate(lower_parts[:-1]) if part in target_aliases), None)
+    other_index = next((index for index, part in enumerate(lower_parts[:-1]) if part in other_aliases), None)
+    if target_index is not None:
+        raw_parts = raw_parts[target_index + 1 :]
+    elif other_index is not None:
+        return None
+
+    if not raw_parts:
+        return None
+    safe_parts = [slugify(part) for part in raw_parts[:-1]]
+    filename = safe_filename(raw_parts[-1])
+    return Path(*safe_parts, filename) if safe_parts else Path(filename)
+
+
+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:
+        return None
+    dst = unique_path(target / relative)
+    resolved_target = target.resolve()
+    resolved_dst = dst.resolve()
+    if resolved_target not in resolved_dst.parents and resolved_dst != resolved_target:
+        raise ValueError(f"archive member escapes target directory: {member_name}")
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    with dst.open("wb") as handle:
+        shutil.copyfileobj(source, handle)
+    return {"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size, "from_archive": True}
+
+
 def dataset_dir(name: str) -> Path:
     return uploads_root() / slugify(name)
 
@@ -283,15 +353,35 @@ async def save_upload(dataset: str, kind: str, files: list[UploadFile]) -> dict:
     saved = []
     for upload in files:
         filename = safe_filename(upload.filename)
-        dst = target / filename
-        if dst.exists():
-            stem = dst.stem
-            suffix = dst.suffix
-            counter = 1
-            while dst.exists():
-                dst = target / f"{stem}_{counter}{suffix}"
-                counter += 1
+        upload.file.seek(0)
+        if is_archive_name(upload.filename):
+            archive_saved = []
+            if filename.lower().endswith(".zip"):
+                with zipfile.ZipFile(upload.file) as archive:
+                    for info in archive.infolist():
+                        if info.is_dir():
+                            continue
+                        with archive.open(info) as source:
+                            item = save_archive_member(target, info.filename, kind, source)
+                            if item:
+                                archive_saved.append(item)
+            else:
+                with tarfile.open(fileobj=upload.file, mode="r:*") as archive:
+                    for member in archive:
+                        if not member.isfile():
+                            continue
+                        source = archive.extractfile(member)
+                        if source is None:
+                            continue
+                        with source:
+                            item = save_archive_member(target, member.name, kind, source)
+                            if item:
+                                archive_saved.append(item)
+            saved.extend(archive_saved)
+            continue
+
+        dst = unique_path(target / filename)
         with dst.open("wb") as handle:
             shutil.copyfileobj(upload.file, handle)
-        saved.append({"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size})
+        saved.append({"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size, "from_archive": False})
     return {"dataset": describe_dataset(safe_name), "saved": saved}

backend/tests/test_dataset_service.py

@@ -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