Support suffixed mask pairing
This commit is contained in:
@@ -22,6 +22,7 @@ ARCHIVE_ALIASES = {
|
|||||||
"masks": {"mask", "masks", "label", "labels", "gt", "annotation", "annotations"},
|
"masks": {"mask", "masks", "label", "labels", "gt", "annotation", "annotations"},
|
||||||
"labels": {"label", "labels", "txt", "annotation", "annotations"},
|
"labels": {"label", "labels", "txt", "annotation", "annotations"},
|
||||||
}
|
}
|
||||||
|
MASK_STEM_SUFFIXES = ("-mask", "_mask", "-label", "_label", "-gt", "_gt", "-seg", "_seg")
|
||||||
|
|
||||||
|
|
||||||
def uploads_root() -> Path:
|
def uploads_root() -> Path:
|
||||||
@@ -44,6 +45,16 @@ def safe_filename(value: str | None) -> str:
|
|||||||
return stem
|
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:
|
def is_archive_name(value: str | None) -> bool:
|
||||||
name = (value or "").lower()
|
name = (value or "").lower()
|
||||||
return any(name.endswith(suffix) for suffix in ARCHIVE_SUFFIXES)
|
return any(name.endswith(suffix) for suffix in ARCHIVE_SUFFIXES)
|
||||||
@@ -88,9 +99,19 @@ def safe_archive_member(raw_name: str, kind: str) -> Path | None:
|
|||||||
return None
|
return None
|
||||||
safe_parts = [slugify(part) for part in raw_parts[:-1]]
|
safe_parts = [slugify(part) for part in raw_parts[:-1]]
|
||||||
filename = safe_filename(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)
|
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:
|
def save_archive_member(target: Path, member_name: str, kind: str, source) -> dict | None:
|
||||||
relative = safe_archive_member(member_name, kind)
|
relative = safe_archive_member(member_name, kind)
|
||||||
if relative is None:
|
if relative is None:
|
||||||
@@ -147,12 +168,13 @@ def _iter_files(root: Path) -> Iterable[Path]:
|
|||||||
return sorted(path for path in root.rglob("*") if path.is_file())
|
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]:
|
def _stem_map(paths: Iterable[Path], exts: set[str] | None = None, normalize=None) -> dict[str, Path]:
|
||||||
result: dict[str, Path] = {}
|
result: dict[str, Path] = {}
|
||||||
for path in paths:
|
for path in paths:
|
||||||
if exts and path.suffix.lower() not in exts:
|
if exts and path.suffix.lower() not in exts:
|
||||||
continue
|
continue
|
||||||
result.setdefault(path.stem, path)
|
stem = normalize(path.stem) if normalize else path.stem
|
||||||
|
result.setdefault(stem, path)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
@@ -234,7 +256,7 @@ def validate_dataset(name: str) -> dict:
|
|||||||
mask_files = list(_iter_files(root / "masks"))
|
mask_files = list(_iter_files(root / "masks"))
|
||||||
images = _stem_map(image_files, IMAGE_EXTS)
|
images = _stem_map(image_files, IMAGE_EXTS)
|
||||||
labels = _stem_map(label_files, LABEL_EXTS)
|
labels = _stem_map(label_files, LABEL_EXTS)
|
||||||
masks = _stem_map(mask_files, IMAGE_EXTS)
|
masks = _stem_map(mask_files, IMAGE_EXTS, normalize_mask_stem)
|
||||||
|
|
||||||
image_stems = set(images)
|
image_stems = set(images)
|
||||||
label_stems = set(labels)
|
label_stems = set(labels)
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ def test_validate_dataset_and_generate_yolo_yaml(tmp_path, monkeypatch):
|
|||||||
image = np.zeros((32, 32, 3), dtype=np.uint8)
|
image = np.zeros((32, 32, 3), dtype=np.uint8)
|
||||||
mask = np.zeros((32, 32), dtype=np.uint8)
|
mask = np.zeros((32, 32), dtype=np.uint8)
|
||||||
cv2.imwrite(str(root / "images" / "sample.png"), image)
|
cv2.imwrite(str(root / "images" / "sample.png"), image)
|
||||||
cv2.imwrite(str(root / "masks" / "sample.png"), mask)
|
cv2.imwrite(str(root / "masks" / "sample-mask.png"), mask)
|
||||||
(root / "labels" / "sample.txt").write_text("0 0.25 0.25 0.75 0.25 0.75 0.75 0.25 0.75\n", encoding="utf-8")
|
(root / "labels" / "sample.txt").write_text("0 0.25 0.25 0.75 0.25 0.75 0.75 0.25 0.75\n", encoding="utf-8")
|
||||||
|
|
||||||
validation = validate_dataset("case_yolo")
|
validation = validate_dataset("case_yolo")
|
||||||
@@ -57,6 +57,7 @@ def test_upload_zip_extracts_matching_dataset_kind(tmp_path, monkeypatch):
|
|||||||
image_zip = io.BytesIO()
|
image_zip = io.BytesIO()
|
||||||
with zipfile.ZipFile(image_zip, "w") as archive:
|
with zipfile.ZipFile(image_zip, "w") as archive:
|
||||||
archive.writestr("bundle/images/sample.png", b"image")
|
archive.writestr("bundle/images/sample.png", b"image")
|
||||||
|
archive.writestr("bundle/images/manifest.csv", b"name\nsample.png\n")
|
||||||
archive.writestr("bundle/masks/sample.png", b"mask")
|
archive.writestr("bundle/masks/sample.png", b"mask")
|
||||||
image_zip.seek(0)
|
image_zip.seek(0)
|
||||||
image_upload = UploadFile(filename="images.zip", file=image_zip)
|
image_upload = UploadFile(filename="images.zip", file=image_zip)
|
||||||
@@ -65,13 +66,13 @@ def test_upload_zip_extracts_matching_dataset_kind(tmp_path, monkeypatch):
|
|||||||
mask_zip = io.BytesIO()
|
mask_zip = io.BytesIO()
|
||||||
with zipfile.ZipFile(mask_zip, "w") as archive:
|
with zipfile.ZipFile(mask_zip, "w") as archive:
|
||||||
archive.writestr("bundle/images/sample.png", b"image")
|
archive.writestr("bundle/images/sample.png", b"image")
|
||||||
archive.writestr("bundle/masks/sample.png", b"mask")
|
archive.writestr("bundle/masks/sample-mask.png", b"mask")
|
||||||
mask_zip.seek(0)
|
mask_zip.seek(0)
|
||||||
mask_upload = UploadFile(filename="masks.zip", file=mask_zip)
|
mask_upload = UploadFile(filename="masks.zip", file=mask_zip)
|
||||||
mask_result = asyncio.run(save_upload("case_zip", "masks", [mask_upload]))
|
mask_result = asyncio.run(save_upload("case_zip", "masks", [mask_upload]))
|
||||||
|
|
||||||
assert [item["relative_path"] for item in image_result["saved"]] == ["var/uploads/datasets/case_zip/images/sample.png"]
|
assert [item["relative_path"] for item in image_result["saved"]] == ["var/uploads/datasets/case_zip/images/sample.png"]
|
||||||
assert [item["relative_path"] for item in mask_result["saved"]] == ["var/uploads/datasets/case_zip/masks/sample.png"]
|
assert [item["relative_path"] for item in mask_result["saved"]] == ["var/uploads/datasets/case_zip/masks/sample-mask.png"]
|
||||||
described = describe_dataset("case_zip")
|
described = describe_dataset("case_zip")
|
||||||
assert described["counts"]["images"] == 1
|
assert described["counts"]["images"] == 1
|
||||||
assert described["counts"]["masks"] == 1
|
assert described["counts"]["masks"] == 1
|
||||||
|
|||||||
Reference in New Issue
Block a user