scripts/prepare_polygon_dataset.py

"""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 = {"bbox", "box_txt", "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("--bbox-source-root", default=None, help="Optional local bbox dataset mirror for 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_bbox_image(record: dict[str, Any], bbox_root: Path | None) -> Path | None:
    if bbox_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 = bbox_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)
    bbox_root = Path(args.bbox_source_root) if args.bbox_source_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_bbox_image(record, bbox_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()