133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import shutil
|
|
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"}
|
|
|
|
|
|
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 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 describe_dataset(name: str) -> dict:
|
|
safe_name = slugify(name)
|
|
root = dataset_dir(safe_name)
|
|
meta = _load_meta(safe_name)
|
|
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, "counts": counts, "samples": samples}
|
|
|
|
|
|
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)
|
|
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
|
|
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})
|
|
return {"dataset": describe_dataset(safe_name), "saved": saved}
|