298 lines
11 KiB
Python
298 lines
11 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"}
|
|
LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}
|
|
|
|
|
|
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 _stem_map(paths: Iterable[Path], exts: set[str] | None = None) -> dict[str, Path]:
|
|
result: dict[str, Path] = {}
|
|
for path in paths:
|
|
if exts and path.suffix.lower() not in exts:
|
|
continue
|
|
result.setdefault(path.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)
|
|
|
|
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)
|
|
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}
|