Redesign dataset workbench management

This commit is contained in:
2026-07-01 11:43:30 +08:00
parent 9d66f65212
commit b1c0c982b8
6 changed files with 740 additions and 75 deletions

View File

@@ -16,6 +16,7 @@ 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"},
@@ -47,7 +48,7 @@ def safe_filename(value: str | None) -> str:
def upload_exts_for_kind(kind: str) -> set[str]:
if kind == "images":
return IMAGE_EXTS
return IMAGE_EXTS | VIDEO_EXTS
if kind == "masks":
return IMAGE_EXTS | LABEL_EXTS
if kind == "labels":
@@ -138,6 +139,8 @@ def metadata_path(name: str) -> Path:
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 = {
@@ -155,6 +158,44 @@ def create_dataset(name: str, description: str = "") -> dict:
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():
@@ -192,6 +233,22 @@ def _image_shape(path: Path) -> dict | None:
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()
@@ -232,20 +289,14 @@ def describe_dataset(name: str) -> dict:
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)
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}
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: