461 lines
18 KiB
Python
461 lines
18 KiB
Python
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"}
|
|
VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv", ".webm", ".m4v"}
|
|
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 | VIDEO_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)
|
|
if root.exists():
|
|
raise ValueError(f"dataset already exists: {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 delete_dataset(name: str) -> dict:
|
|
safe_name = slugify(name)
|
|
root = dataset_dir(safe_name)
|
|
if not root.exists() or not root.is_dir():
|
|
raise ValueError(f"dataset not found: {safe_name}")
|
|
shutil.rmtree(root)
|
|
return {"deleted": True, "name": safe_name}
|
|
|
|
|
|
def copy_dataset(source_name: str, target_name: str, description: str | None = None) -> dict:
|
|
source_safe = slugify(source_name)
|
|
target_safe = slugify(target_name)
|
|
source = dataset_dir(source_safe)
|
|
target = dataset_dir(target_safe)
|
|
if not source.exists() or not source.is_dir():
|
|
raise ValueError(f"dataset not found: {source_safe}")
|
|
if target.exists():
|
|
raise ValueError(f"dataset already exists: {target_safe}")
|
|
shutil.copytree(source, target)
|
|
meta = _load_meta(source_safe)
|
|
meta.update(
|
|
{
|
|
"name": target_safe,
|
|
"description": description if description is not None else f"Copy of {source_safe}",
|
|
"created_at": datetime.now(timezone.utc).isoformat(),
|
|
"copied_from": source_safe,
|
|
"root": str(target.relative_to(settings.project_root)),
|
|
"layout": {
|
|
"images": str((target / "images").relative_to(settings.project_root)),
|
|
"labels": str((target / "labels").relative_to(settings.project_root)),
|
|
"masks": str((target / "masks").relative_to(settings.project_root)),
|
|
},
|
|
}
|
|
)
|
|
metadata_path(target_safe).write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return describe_dataset(target_safe)
|
|
|
|
|
|
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 _file_info(path: Path, kind: str) -> dict:
|
|
suffix = path.suffix.lower()
|
|
stem = path.stem
|
|
return {
|
|
"name": path.name,
|
|
"stem": stem,
|
|
"pair_stem": normalize_mask_stem(stem) if kind == "masks" else stem,
|
|
"path": str(path.resolve()),
|
|
"relative_path": str(path.resolve().relative_to(settings.project_root)),
|
|
"size": path.stat().st_size,
|
|
"ext": suffix,
|
|
"previewable": suffix in IMAGE_EXTS,
|
|
"media_type": "image" if suffix in IMAGE_EXTS else "video" if suffix in VIDEO_EXTS else "file",
|
|
}
|
|
|
|
|
|
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 = {}
|
|
files_by_kind = {}
|
|
for kind in sorted(DATASET_KINDS):
|
|
files = list(_iter_files(root / kind))
|
|
counts[kind] = len(files)
|
|
file_items = [_file_info(path, kind) for path in files[:2000]]
|
|
files_by_kind[kind] = file_items
|
|
samples[kind] = file_items[:80]
|
|
return {**meta, "absolute_layout": absolute_layout, "counts": counts, "samples": samples, "files": files_by_kind}
|
|
|
|
|
|
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}
|