scripts/import_patch_polygon_dataset.py

"""Import patch folders into the project's normalized polygon dataset layout."""

from __future__ import annotations

import argparse
import csv
import json
import random
import re
import shutil
from collections import Counter, defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw


ELEMENT_MAP = {
    "海上养殖区": {
        "id": "aquaculture",
        "name": "海上养殖区",
        "task_type": "polygon_extraction",
    },
    "港口岸线": {
        "id": "port_shoreline",
        "name": "港口岸线",
        "task_type": "polygon_extraction",
    },
    "粉砂淤泥质岸线": {
        "id": "silty_muddy_shoreline",
        "name": "粉砂淤泥质岸线",
        "task_type": "polygon_extraction",
    },
}

IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
KEY_SUFFIX = re.compile(r"_(Orig|True|False|TrueColor|FalseColor|Binary|Label)_[^_]+$")


@dataclass
class ImportStats:
    scanned: int = 0
    accepted: int = 0
    rejected: int = 0
    inference_assets: int = 0
    missing_train_image: int = 0
    missing_mask: int = 0
    missing_geojson: int = 0
    empty_mask: int = 0
    invalid_geojson: int = 0
    size_mismatch: int = 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-root", required=True)
    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=20260704)
    parser.add_argument("--preview-limit", type=int, default=36)
    return parser.parse_args()


def read_geojson(path: Path) -> dict[str, Any] | None:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    if data.get("type") != "FeatureCollection":
        return None
    features = data.get("features")
    if not isinstance(features, list) or len(features) == 0:
        return None
    for feature in features:
        geom = feature.get("geometry") if isinstance(feature, dict) else None
        if isinstance(geom, dict) and geom.get("type") in {"Polygon", "MultiPolygon", "LineString", "MultiLineString"}:
            return data
    return None


def file_key(path: Path) -> str:
    return KEY_SUFFIX.sub("", path.stem)


def index_files(directory: Path) -> dict[str, Path]:
    if not directory.exists():
        return {}
    return {file_key(path): path for path in directory.iterdir() if path.is_file()}


def safe_name(text: str) -> str:
    return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in text)


def mask_stats(mask_path: Path) -> tuple[tuple[int, int] | None, int, float]:
    try:
        mask = Image.open(mask_path).convert("L")
    except Exception:
        return None, 0, 0.0
    extrema = mask.getextrema()
    if extrema is None or extrema[1] == 0:
        return mask.size, 0, 0.0
    binary = mask.point(lambda v: 255 if v > 0 else 0)
    foreground = 0
    hist = binary.histogram()
    if len(hist) > 255:
        foreground = hist[255]
    ratio = foreground / max(binary.size[0] * binary.size[1], 1)
    return binary.size, foreground, ratio


def image_size(path: Path) -> tuple[int, int] | None:
    try:
        return Image.open(path).size
    except Exception:
        return None


def discover_groups(source_root: Path) -> list[tuple[str, str, int, Path]]:
    groups = []
    for element_dir in sorted(source_root.iterdir()):
        if not element_dir.is_dir():
            continue
        for satellite_dir in sorted(element_dir.iterdir()):
            if not satellite_dir.is_dir():
                continue
            for size_dir in sorted(satellite_dir.glob("Size_*")):
                if not size_dir.is_dir():
                    continue
                try:
                    patch_size = int(size_dir.name.split("_", 1)[1])
                except Exception:
                    continue
                groups.append((element_dir.name, satellite_dir.name, patch_size, size_dir))
    return groups


def build_records(source_root: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]], ImportStats]:
    accepted: list[dict[str, Any]] = []
    rejected: list[dict[str, Any]] = []
    stats = ImportStats()
    for element_name, satellite, patch_size, size_dir in discover_groups(source_root):
        element = ELEMENT_MAP.get(element_name, {"id": safe_name(element_name), "name": element_name, "task_type": "polygon_extraction"})
        true_images = index_files(size_dir / "Image_TrueColor")
        false_images = index_files(size_dir / "Image_FalseColor")
        orig_images = index_files(size_dir / "Image_Orig")
        masks = index_files(size_dir / "Label_Binary")
        geojsons = index_files(size_dir / "Label_GeoJSON")
        keys = sorted(set(true_images) | set(false_images) | set(orig_images) | set(masks) | set(geojsons))
        for key in keys:
            stats.scanned += 1
            true_path = true_images.get(key)
            false_path = false_images.get(key)
            orig_path = orig_images.get(key)
            image_path = true_path or false_path
            mask_path = masks.get(key)
            geojson_path = geojsons.get(key)
            reject_reason = None
            if image_path is None:
                reject_reason = "missing_true_or_false_color_image"
                stats.missing_train_image += 1
            elif mask_path is None:
                reject_reason = "missing_binary_mask"
                stats.missing_mask += 1
            elif geojson_path is None:
                reject_reason = "missing_geojson_polygon"
                stats.missing_geojson += 1
            else:
                img_size = image_size(image_path)
                m_size, foreground, fg_ratio = mask_stats(mask_path)
                geojson = read_geojson(geojson_path)
                if m_size is None:
                    reject_reason = "unreadable_binary_mask"
                    stats.missing_mask += 1
                elif foreground == 0:
                    reject_reason = "empty_binary_mask"
                    stats.empty_mask += 1
                elif img_size is not None and img_size != m_size:
                    reject_reason = "image_mask_size_mismatch"
                    stats.size_mismatch += 1
                elif geojson is None:
                    reject_reason = "invalid_or_empty_geojson"
                    stats.invalid_geojson += 1
                else:
                    sample_id = safe_name(f"{element['id']}_{satellite}_{patch_size}_{key}")
                    accepted.append(
                        {
                            "sample_id": sample_id,
                            "source_key": key,
                            "element": element["id"],
                            "element_name": element["name"],
                            "task_type": element["task_type"],
                            "image_path": str(image_path),
                            "mask_path": str(mask_path),
                            "annotation_path": str(geojson_path),
                            "annotation_format": "geojson_polygon",
                            "label_encoding": {"0": "background", "1": element["id"]},
                            "satellite": satellite,
                            "sensor": infer_sensor(key),
                            "resolution_m": infer_resolution_m(satellite),
                            "patch_size": patch_size,
                            "bands": ["red", "green", "blue"],
                            "band_count": 3,
                            "dtype": "uint8",
                            "fusion": {
                                "state": infer_fusion_state(satellite),
                                "method": "unknown_patch_product",
                                "sources": [
                                    {"role": "true_color", "path": str(true_path) if true_path else None, "resolution_m": infer_resolution_m(satellite)},
                                    {"role": "false_color", "path": str(false_path) if false_path else None, "resolution_m": infer_resolution_m(satellite)},
                                    {"role": "original_patch", "path": str(orig_path) if orig_path else None, "resolution_m": infer_resolution_m(satellite)},
                                ],
                                "target_resolution_m": infer_resolution_m(satellite),
                                "native_multispectral_resolution_m": None,
                                "persisted": True,
                                "reproducible": False,
                                "spectral_preservation": "unknown",
                            },
                            "source_project": "BaiduNetdisk_Patches",
                            "source_dataset": str(source_root),
                            "quality_flags": [
                                "accepted_for_polygon_training",
                                "paired_image_mask_geojson",
                                f"foreground_ratio:{fg_ratio:.6f}",
                            ],
                            "quality_score": 0.96,
                        }
                    )
                    stats.accepted += 1
                    continue
            item = {
                "source_key": key,
                "element": element["id"],
                "element_name": element["name"],
                "satellite": satellite,
                "patch_size": patch_size,
                "image_path": str(image_path) if image_path else None,
                "orig_image_path": str(orig_path) if orig_path else None,
                "mask_path": str(mask_path) if mask_path else None,
                "annotation_path": str(geojson_path) if geojson_path else None,
                "polygon_import_status": "rejected",
                "reject_reason": reject_reason,
            }
            if image_path and not mask_path and not geojson_path:
                item["polygon_import_status"] = "inference_asset"
                stats.inference_assets += 1
            else:
                stats.rejected += 1
            rejected.append(item)
    return accepted, rejected, stats


def infer_sensor(key: str) -> str | None:
    for token in ("PMS1", "PMS2", "PMS", "MUX"):
        if token in key:
            return token
    return None


def infer_resolution_m(satellite: str) -> float | None:
    return {"GF1": 2.0, "GF2": 1.0, "GF6": 2.0}.get(satellite.upper())


def infer_fusion_state(satellite: str) -> str:
    return "fused_product" if satellite.upper() in {"GF1", "GF2", "GF6"} else "unknown"


def split_records(records: list[dict[str, Any]], train_ratio: float, val_ratio: float, seed: int) -> dict[str, list[dict[str, Any]]]:
    grouped: dict[tuple[str, str, int], list[dict[str, Any]]] = defaultdict(list)
    for record in records:
        grouped[(record["element"], record["satellite"], int(record["patch_size"]))].append(record)
    rng = random.Random(seed)
    splits = {"train": [], "val": [], "test": []}
    for rows in grouped.values():
        rng.shuffle(rows)
        n = len(rows)
        n_train = int(n * train_ratio)
        n_val = int(n * val_ratio)
        if n >= 3:
            n_train = max(1, min(n_train, n - 2))
            n_val = max(1, min(n_val, n - n_train - 1))
        splits["train"].extend(rows[:n_train])
        splits["val"].extend(rows[n_train : n_train + n_val])
        splits["test"].extend(rows[n_train + n_val :])
    return splits


def copy_record(record: dict[str, Any], split: str, output_root: Path) -> dict[str, Any]:
    image_src = Path(record["image_path"])
    mask_src = Path(record["mask_path"])
    geojson_src = Path(record["annotation_path"])
    sample_id = record["sample_id"]
    image_dst = output_root / "images" / split / f"{sample_id}.jpg"
    mask_dst = output_root / "masks" / split / f"{sample_id}.png"
    annotation_dst = output_root / "annotations" / split / f"{sample_id}.geojson"
    image_dst.parent.mkdir(parents=True, exist_ok=True)
    mask_dst.parent.mkdir(parents=True, exist_ok=True)
    annotation_dst.parent.mkdir(parents=True, exist_ok=True)
    Image.open(image_src).convert("RGB").save(image_dst, quality=95)
    mask = Image.open(mask_src).convert("L").point(lambda v: 255 if v > 0 else 0)
    mask.save(mask_dst)
    shutil.copy2(geojson_src, annotation_dst)
    copied = dict(record)
    copied.update({"split": split, "image_path": str(image_dst), "mask_path": str(mask_dst), "annotation_path": str(annotation_dst)})
    return copied


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 write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    keys = sorted({key for row in rows for key in row.keys()})
    with path.open("w", newline="", encoding="utf-8-sig") as fp:
        writer = csv.DictWriter(fp, fieldnames=keys)
        writer.writeheader()
        writer.writerows(rows)


def make_preview(rows: list[dict[str, Any]], output_path: Path, limit: int) -> None:
    rows = rows[:limit]
    if not rows:
        return
    tile = 256
    cols = 6
    rows_n = (len(rows) + cols - 1) // cols
    canvas = Image.new("RGB", (cols * tile, rows_n * tile), (30, 30, 30))
    for idx, row in enumerate(rows):
        image = Image.open(row["image_path"]).convert("RGB").resize((tile, tile))
        mask = Image.open(row["mask_path"]).convert("L").resize((tile, tile))
        overlay = Image.new("RGBA", (tile, tile), (255, 60, 40, 0))
        overlay.putalpha(mask.point(lambda v: 100 if v > 0 else 0))
        image = Image.alpha_composite(image.convert("RGBA"), overlay).convert("RGB")
        draw = ImageDraw.Draw(image)
        draw.text((6, 6), f"{row['element']} {row['satellite']} S{row['patch_size']}", fill=(255, 255, 255))
        x = (idx % cols) * tile
        y = (idx // cols) * tile
        canvas.paste(image, (x, y))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(output_path, quality=92)


def main() -> None:
    args = parse_args()
    source_root = Path(args.source_root)
    output_root = Path(args.output_root)
    if output_root.exists():
        shutil.rmtree(output_root)
    output_root.mkdir(parents=True, exist_ok=True)
    accepted, rejected, stats = build_records(source_root)
    splits = split_records(accepted, args.train_ratio, args.val_ratio, args.seed)
    copied: list[dict[str, Any]] = []
    for split, rows in splits.items():
        for row in rows:
            copied.append(copy_record(row, split, output_root))

    manifests = output_root / "manifests"
    write_jsonl(manifests / "accepted_polygon_samples.jsonl", copied)
    write_jsonl(manifests / "rejected_samples.jsonl", rejected)
    write_csv(output_root / "reports" / "rejected_samples.csv", rejected)
    write_csv(output_root / "reports" / "accepted_samples.csv", copied)
    make_preview(copied, output_root / "previews" / "mask_overlay_contact_sheet.jpg", args.preview_limit)

    by_element = Counter(row["element"] for row in copied)
    by_satellite = Counter(row["satellite"] for row in copied)
    by_patch_size = Counter(str(row["patch_size"]) for row in copied)
    summary = {
        **asdict(stats),
        "output_root": str(output_root),
        "splits": {split: len(rows) for split, rows in splits.items()},
        "accepted_by_element": dict(by_element),
        "accepted_by_satellite": dict(by_satellite),
        "accepted_by_patch_size": dict(by_patch_size),
        "dataset_layout": "images/{split}, masks/{split}, annotations/{split}, manifests, reports, previews",
        "quality_policy": "Only paired image + non-empty binary mask + valid GeoJSON polygon samples are accepted.",
        "source_root": str(source_root),
    }
    (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()