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}