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 from fastapi import UploadFile 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"}, } MASK_STEM_SUFFIXES = ("-mask", "_mask", "-label", "_label", "-gt", "_gt", "-seg", "_seg") def uploads_root() -> Path: root = settings.project_root / "var" / "uploads" / "datasets" root.mkdir(parents=True, exist_ok=True) return root def slugify(value: str) -> str: text = re.sub(r"[^A-Za-z0-9_.\-\u4e00-\u9fff]+", "_", value.strip()) return text.strip("._") or "dataset" def safe_filename(value: str | None) -> str: original = Path(value or "upload.bin").name suffix = Path(original).suffix.lower() stem = slugify(Path(original).stem or "upload") if suffix and re.fullmatch(r"\.[A-Za-z0-9]{1,12}", suffix): return f"{stem}{suffix}" 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) 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]) 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: 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) def metadata_path(name: str) -> Path: return dataset_dir(name) / "dataset.json" def create_dataset(name: str, description: str = "") -> dict: safe_name = slugify(name) root = dataset_dir(safe_name) for kind in DATASET_KINDS: (root / kind).mkdir(parents=True, exist_ok=True) meta = { "name": safe_name, "description": description, "created_at": datetime.now(timezone.utc).isoformat(), "root": str(root.relative_to(settings.project_root)), "layout": { "images": str((root / "images").relative_to(settings.project_root)), "labels": str((root / "labels").relative_to(settings.project_root)), "masks": str((root / "masks").relative_to(settings.project_root)), }, } metadata_path(safe_name).write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8") return describe_dataset(safe_name) def _load_meta(name: str) -> dict: path = metadata_path(name) if path.exists(): return json.loads(path.read_text(encoding="utf-8")) return {"name": slugify(name), "description": "", "root": str(dataset_dir(name).relative_to(settings.project_root))} def _iter_files(root: Path) -> Iterable[Path]: if not root.exists(): return [] return sorted(path for path in root.rglob("*") if path.is_file()) 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 stem = normalize(path.stem) if normalize else path.stem result.setdefault(stem, path) return result def _image_shape(path: Path) -> dict | None: try: import cv2 image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) if image is None: return None height, width = image.shape[:2] channels = 1 if image.ndim == 2 else image.shape[2] return {"width": int(width), "height": int(height), "channels": int(channels)} except Exception: return None def _validate_yolo_txt(path: Path) -> dict: errors = [] classes: set[int] = set() lines = path.read_text(encoding="utf-8", errors="replace").splitlines() annotation_count = 0 for line_number, line in enumerate(lines, 1): raw = line.strip() if not raw: continue parts = raw.split() annotation_count += 1 if len(parts) < 5: errors.append(f"{path.name}:{line_number} has fewer than 5 tokens") continue try: class_id = int(float(parts[0])) coords = [float(item) for item in parts[1:]] except ValueError: errors.append(f"{path.name}:{line_number} contains non-numeric values") continue if class_id < 0: errors.append(f"{path.name}:{line_number} has negative class id") classes.add(class_id) if len(coords) not in {4} and len(coords) < 6: errors.append(f"{path.name}:{line_number} has too few segmentation coordinates") if len(coords) != 4 and len(coords) % 2 != 0: errors.append(f"{path.name}:{line_number} has an odd number of polygon coordinates") out_of_range = [value for value in coords if value < 0 or value > 1] if out_of_range: errors.append(f"{path.name}:{line_number} has coordinates outside 0..1") return {"annotations": annotation_count, "classes": sorted(classes), "errors": errors} def describe_dataset(name: str) -> dict: safe_name = slugify(name) root = dataset_dir(safe_name) meta = _load_meta(safe_name) absolute_layout = {kind: str((root / kind).resolve()) for kind in DATASET_KINDS} counts = {} samples = {} for kind in sorted(DATASET_KINDS): files = list(_iter_files(root / kind)) counts[kind] = len(files) samples[kind] = [ { "name": path.name, "path": str(path.resolve()), "relative_path": str(path.resolve().relative_to(settings.project_root)), "size": path.stat().st_size, "previewable": path.suffix.lower() in IMAGE_EXTS, } for path in files[:80] ] return {**meta, "absolute_layout": absolute_layout, "counts": counts, "samples": samples} def validate_dataset(name: str) -> dict: safe_name = slugify(name) root = dataset_dir(safe_name) image_files = list(_iter_files(root / "images")) label_files = list(_iter_files(root / "labels")) 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, normalize_mask_stem) image_stems = set(images) label_stems = set(labels) mask_stems = set(masks) paired_label_stems = sorted(image_stems & label_stems) paired_mask_stems = sorted(image_stems & mask_stems) yolo_errors = [] class_ids: set[int] = set() annotation_count = 0 for label in labels.values(): if label.suffix.lower() != ".txt": continue detail = _validate_yolo_txt(label) yolo_errors.extend(detail["errors"]) class_ids.update(detail["classes"]) annotation_count += detail["annotations"] shape_mismatches = [] sample_shapes = [] for stem in paired_mask_stems[:80]: image_shape = _image_shape(images[stem]) mask_shape = _image_shape(masks[stem]) if image_shape and len(sample_shapes) < 8: sample_shapes.append({"name": images[stem].name, **image_shape}) if image_shape and mask_shape and (image_shape["width"], image_shape["height"]) != (mask_shape["width"], mask_shape["height"]): shape_mismatches.append({"stem": stem, "image": image_shape, "mask": mask_shape}) if not sample_shapes: for image in list(images.values())[:8]: image_shape = _image_shape(image) if image_shape: sample_shapes.append({"name": image.name, **image_shape}) checks = [ {"name": "has_images", "passed": len(images) > 0, "count": len(images)}, {"name": "has_labels_or_masks", "passed": len(labels) > 0 or len(masks) > 0, "labels": len(labels), "masks": len(masks)}, {"name": "image_label_pairs", "passed": len(label_stems) == 0 or len(paired_label_stems) > 0, "count": len(paired_label_stems)}, {"name": "image_mask_pairs", "passed": len(mask_stems) == 0 or len(paired_mask_stems) > 0, "count": len(paired_mask_stems)}, {"name": "yolo_txt_valid", "passed": not yolo_errors, "errors": yolo_errors[:20]}, {"name": "mask_shapes_match", "passed": not shape_mismatches, "errors": shape_mismatches[:20]}, ] yolo_ready = len(images) > 0 and len(paired_label_stems) > 0 and not yolo_errors mask_ready = len(images) > 0 and len(paired_mask_stems) > 0 and not shape_mismatches return { "dataset": safe_name, "root": str(root.resolve()), "counts": {"images": len(images), "labels": len(labels), "masks": len(masks), "annotations": annotation_count}, "pairs": { "image_label": len(paired_label_stems), "image_mask": len(paired_mask_stems), "images_without_labels": sorted(image_stems - label_stems)[:50], "labels_without_images": sorted(label_stems - image_stems)[:50], "images_without_masks": sorted(image_stems - mask_stems)[:50], "masks_without_images": sorted(mask_stems - image_stems)[:50], }, "classes": sorted(class_ids), "sample_shapes": sample_shapes, "checks": checks, "ready": {"yolo": yolo_ready, "mask": mask_ready, "any": yolo_ready or mask_ready}, } def generate_yolo_dataset_yaml(name: str, class_names: list[str] | None = None) -> dict: validation = validate_dataset(name) if not validation["ready"]["yolo"]: raise ValueError("dataset is not YOLO-ready; upload matching images and valid txt labels first") safe_name = slugify(name) root = dataset_dir(safe_name) classes = validation["classes"] or [0] class_count = max(classes) + 1 names = list(class_names or []) if len(names) < class_count: names.extend(f"class_{index}" for index in range(len(names), class_count)) yaml_path = root / "dataset.yaml" names_block = "\n".join(f" {index}: {label}" for index, label in enumerate(names)) yaml_text = "\n".join( [ f"path: {root.resolve()}", "train: images", "val: images", f"nc: {len(names)}", "names:", names_block, "", ] ) yaml_path.write_text(yaml_text, encoding="utf-8") return { "dataset": safe_name, "path": str(yaml_path.resolve()), "relative_path": str(yaml_path.resolve().relative_to(settings.project_root)), "classes": classes, "names": names, "content": yaml_text, "validation": validation, } def list_uploaded_datasets() -> list[dict]: root = uploads_root() datasets = [] for item in sorted(root.iterdir()): if item.is_dir(): datasets.append(describe_dataset(item.name)) return datasets async def save_upload(dataset: str, kind: str, files: list[UploadFile]) -> dict: if kind not in DATASET_KINDS: raise ValueError(f"unsupported dataset file kind: {kind}") safe_name = slugify(dataset) if not metadata_path(safe_name).exists(): create_dataset(safe_name) target = dataset_dir(safe_name) / kind target.mkdir(parents=True, exist_ok=True) saved = [] for upload in files: filename = safe_filename(upload.filename) 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, "from_archive": False}) return {"dataset": describe_dataset(safe_name), "saved": saved}