Sync strict polygon dataset importer

scripts/prepare_polygon_dataset.py

@@ -0,0 +1,276 @@
+"""Build the normalized polygon dataset used by SAMPoly-style training.
+
+The importer is intentionally strict: bbox-only annotations are rejected because
+they cannot supervise true polygon boundaries or vertices.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+from PIL import Image
+
+
+IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
+MASK_SUFFIXES = {".png", ".tif", ".tiff", ".jpg", ".jpeg"}
+POLYGON_FORMATS = {"coco_polygon", "coco_segmentation", "geojson", "shp", "mask", "binary_mask", "semantic_mask"}
+BBOX_FORMATS = {"yolo_bbox_txt", "bbox", "coco_bbox", "voc_bbox"}
+
+
+@dataclass
+class ImportStats:
+    scanned: int = 0
+    accepted: int = 0
+    rejected: int = 0
+    accepted_masks: int = 0
+    accepted_polygons: int = 0
+    rejected_bbox_only: int = 0
+    rejected_missing_image: int = 0
+    rejected_missing_label: int = 0
+    rejected_unknown_format: int = 0
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--manifest", required=True, help="JSONL manifest with standardized sample records.")
+    parser.add_argument("--local-yolo-root", default=None, help="Optional local YOLO mirror for bbox rejection auditing.")
+    parser.add_argument("--extra-source-root", action="append", default=[], help="Local source roots to scan for mask/polygon datasets.")
+    parser.add_argument("--output-root", required=True)
+    parser.add_argument("--train-ratio", type=float, default=0.8)
+    parser.add_argument("--val-ratio", type=float, default=0.1)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--min-quality-score", type=float, default=0.9)
+    parser.add_argument("--element", default=None)
+    return parser.parse_args()
+
+
+def read_jsonl(path: Path) -> list[dict[str, Any]]:
+    rows = []
+    if not path.exists():
+        return rows
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if line.strip():
+            rows.append(json.loads(line))
+    return rows
+
+
+def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text("\n".join(json.dumps(row, ensure_ascii=False) for row in rows) + ("\n" if rows else ""), encoding="utf-8")
+
+
+def safe_name(sample_id: str, fallback: str) -> str:
+    raw = sample_id or Path(fallback).stem
+    return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in raw)
+
+
+def local_path_from_record(record: dict[str, Any], key: str) -> Path | None:
+    value = record.get(key)
+    if not value or not isinstance(value, str):
+        return None
+    if value.startswith("hf://"):
+        return None
+    path = Path(value)
+    return path if path.exists() else None
+
+
+def find_local_yolo_image(record: dict[str, Any], yolo_root: Path | None) -> Path | None:
+    if yolo_root is None:
+        return None
+    source = str(record.get("image_path") or "")
+    stem = Path(source).stem.lower()
+    for split in ("train", "val", "test"):
+        image_dir = yolo_root / "images" / split
+        if not image_dir.exists():
+            continue
+        for path in image_dir.iterdir():
+            if path.suffix.lower() in IMAGE_SUFFIXES and path.stem.lower().endswith(stem):
+                return path
+    return None
+
+
+def mask_has_foreground(path: Path) -> bool:
+    try:
+        img = Image.open(path).convert("L")
+        extrema = img.getextrema()
+        return bool(extrema and extrema[1] > 0)
+    except Exception:
+        return False
+
+
+def find_extra_samples(root: Path, min_quality: float, element: str | None) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for image_dir in root.rglob("images"):
+        if not image_dir.is_dir():
+            continue
+        split = image_dir.parent.name if image_dir.parent.name in {"train", "val", "test"} else None
+        mask_dir_candidates = [
+            image_dir.parent / "masks",
+            image_dir.parent.parent / "masks" / image_dir.name,
+            image_dir.parent.parent / "masks" / (split or ""),
+        ]
+        for image_path in image_dir.iterdir():
+            if image_path.suffix.lower() not in IMAGE_SUFFIXES:
+                continue
+            mask_path = None
+            for mask_dir in mask_dir_candidates:
+                if not mask_dir.exists():
+                    continue
+                for suffix in MASK_SUFFIXES:
+                    candidate = mask_dir / f"{image_path.stem}{suffix}"
+                    if candidate.exists():
+                        mask_path = candidate
+                        break
+                if mask_path:
+                    break
+            if not mask_path or not mask_has_foreground(mask_path):
+                continue
+            rows.append(
+                {
+                    "sample_id": f"local_{safe_name(image_path.stem, image_path.name)}",
+                    "element": element or "unknown",
+                    "task_type": "polygon_extraction",
+                    "image_path": str(image_path),
+                    "mask_path": str(mask_path),
+                    "annotation_path": str(mask_path),
+                    "annotation_format": "binary_mask",
+                    "quality_score": max(min_quality, 0.95),
+                    "quality_flags": ["accepted", "local_mask_pair", "polygon_trainable"],
+                    "split": split,
+                }
+            )
+    return rows
+
+
+def split_rows(rows: list[dict[str, Any]], train_ratio: float, val_ratio: float, seed: int) -> dict[str, list[dict[str, Any]]]:
+    grouped = {"train": [], "val": [], "test": []}
+    presplit = [row for row in rows if row.get("split") in grouped]
+    unsplit = [row for row in rows if row.get("split") not in grouped]
+    for row in presplit:
+        grouped[str(row["split"])].append(row)
+    random.Random(seed).shuffle(unsplit)
+    n = len(unsplit)
+    n_train = int(n * train_ratio)
+    n_val = int(n * val_ratio)
+    grouped["train"].extend(unsplit[:n_train])
+    grouped["val"].extend(unsplit[n_train : n_train + n_val])
+    grouped["test"].extend(unsplit[n_train + n_val :])
+    return grouped
+
+
+def copy_sample(row: dict[str, Any], split: str, output_root: Path) -> dict[str, Any]:
+    image_path = Path(str(row["image_path"]))
+    mask_path = Path(str(row.get("mask_path") or row.get("annotation_path")))
+    name = safe_name(str(row.get("sample_id") or image_path.stem), image_path.name)
+    image_out = output_root / "images" / split / f"{name}{image_path.suffix.lower()}"
+    mask_out = output_root / "masks" / split / f"{name}.png"
+    image_out.parent.mkdir(parents=True, exist_ok=True)
+    mask_out.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(image_path, image_out)
+    Image.open(mask_path).convert("L").save(mask_out)
+    copied = dict(row)
+    copied.update(
+        {
+            "sample_id": name,
+            "split": split,
+            "image_path": str(image_out),
+            "mask_path": str(mask_out),
+            "annotation_path": str(mask_out),
+            "annotation_format": "binary_mask",
+            "task_type": "polygon_extraction",
+            "quality_flags": sorted(set(row.get("quality_flags", []) + ["accepted_for_polygon_training"])),
+        }
+    )
+    return copied
+
+
+def main() -> None:
+    args = parse_args()
+    manifest = Path(args.manifest)
+    output_root = Path(args.output_root)
+    output_root.mkdir(parents=True, exist_ok=True)
+    yolo_root = Path(args.local_yolo_root) if args.local_yolo_root else None
+    stats = ImportStats()
+    accepted: list[dict[str, Any]] = []
+    rejected: list[dict[str, Any]] = []
+
+    records = read_jsonl(manifest)
+    for root in args.extra_source_root:
+        records.extend(find_extra_samples(Path(root), args.min_quality_score, args.element))
+
+    for record in records:
+        stats.scanned += 1
+        if args.element and record.get("element") != args.element:
+            continue
+        quality = float(record.get("quality_score") or 0.0)
+        fmt = str(record.get("annotation_format") or "").lower()
+        image_path = local_path_from_record(record, "image_path") or find_local_yolo_image(record, yolo_root)
+        label_path = local_path_from_record(record, "mask_path") or local_path_from_record(record, "annotation_path")
+
+        reject_reason = None
+        if quality < args.min_quality_score:
+            reject_reason = "quality_below_threshold"
+        elif fmt in BBOX_FORMATS:
+            reject_reason = "bbox_only_not_polygon_trainable"
+            stats.rejected_bbox_only += 1
+        elif fmt not in POLYGON_FORMATS:
+            reject_reason = "unknown_or_unsupported_annotation_format"
+            stats.rejected_unknown_format += 1
+        elif image_path is None:
+            reject_reason = "missing_local_image"
+            stats.rejected_missing_image += 1
+        elif label_path is None or not label_path.exists():
+            reject_reason = "missing_local_mask_or_polygon"
+            stats.rejected_missing_label += 1
+        elif fmt in {"mask", "binary_mask", "semantic_mask"} and not mask_has_foreground(label_path):
+            reject_reason = "empty_or_invalid_mask"
+
+        if reject_reason:
+            item = dict(record)
+            item["polygon_import_status"] = "rejected"
+            item["reject_reason"] = reject_reason
+            if image_path:
+                item["local_image_path"] = str(image_path)
+            rejected.append(item)
+            stats.rejected += 1
+            continue
+
+        item = dict(record)
+        item["image_path"] = str(image_path)
+        item["mask_path"] = str(label_path)
+        item["annotation_path"] = str(label_path)
+        item["polygon_import_status"] = "accepted"
+        accepted.append(item)
+        stats.accepted += 1
+        if fmt in {"mask", "binary_mask", "semantic_mask"}:
+            stats.accepted_masks += 1
+        else:
+            stats.accepted_polygons += 1
+
+    grouped = split_rows(accepted, args.train_ratio, args.val_ratio, args.seed)
+    copied_rows = []
+    for split, rows in grouped.items():
+        for row in rows:
+            copied_rows.append(copy_sample(row, split, output_root))
+
+    write_jsonl(output_root / "manifests" / "accepted_polygon_samples.jsonl", copied_rows)
+    write_jsonl(output_root / "manifests" / "rejected_polygon_samples.jsonl", rejected)
+    summary = {
+        **asdict(stats),
+        "output_root": str(output_root),
+        "splits": {split: len(rows) for split, rows in grouped.items()},
+        "quality_policy": "Only mask or polygon annotations are accepted for SAMPoly-style polygon training; bbox-only samples are rejected.",
+        "source_manifest": str(manifest),
+    }
+    (output_root / "dataset_card.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
+    print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True)
+
+
+if __name__ == "__main__":
+    main()