scripts/build_unlabeled_mim_dataset.py

"""Build a manifest-driven unlabeled MIM dataset from remote-sensing images."""

from __future__ import annotations

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

import numpy as np


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}


@dataclass
class Stats:
    scanned_files: int = 0
    accepted_files: int = 0
    rejected_files: int = 0
    accepted_windows: int = 0
    unreadable: int = 0
    too_small: int = 0
    too_black: int = 0
    too_white: int = 0
    low_texture: int = 0
    invalid_shape: int = 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--roots", nargs="+", required=True)
    parser.add_argument("--output-root", required=True)
    parser.add_argument("--patch-size", type=int, default=512)
    parser.add_argument("--stride", type=int, default=512)
    parser.add_argument("--max-files-per-root", type=int, default=30000)
    parser.add_argument("--max-total-windows", type=int, default=120000)
    parser.add_argument("--max-windows-per-image", type=int, default=16)
    parser.add_argument("--progress-every", type=int, default=500)
    parser.add_argument("--min-valid-ratio", type=float, default=0.70)
    parser.add_argument("--max-black-ratio", type=float, default=0.45)
    parser.add_argument("--max-white-ratio", type=float, default=0.65)
    parser.add_argument("--min-std", type=float, default=4.0)
    parser.add_argument("--seed", type=int, default=20260705)
    parser.add_argument("--split-train", type=float, default=0.95)
    return parser.parse_args()


def infer_satellite(text: str) -> str | None:
    upper = text.upper()
    for token in ("GF7", "GF6", "GF5", "GF4", "GF3", "GF2", "GF1", "SENTINEL2", "SENTINEL-2", "LANDSAT8", "LANDSAT9"):
        if token in upper:
            return token.replace("SENTINEL-2", "SENTINEL2")
    match = re.search(r"\bS2[AB]?\b", upper)
    if match:
        return "SENTINEL2"
    return None


def infer_sensor(text: str) -> str | None:
    upper = text.upper()
    for token in ("PMS1", "PMS2", "PMS", "MUX", "PAN", "MSI", "OLI", "SAR", "WFV"):
        if token in upper:
            return token
    return None


def infer_resolution_m(satellite: str | None, sensor: str | None) -> float | None:
    if satellite == "GF2" and sensor in {"PMS1", "PMS2", "PMS"}:
        return 1.0
    if satellite in {"GF1", "GF6"} and sensor in {"PMS1", "PMS2", "PMS"}:
        return 2.0
    if satellite == "SENTINEL2" or sensor == "MSI":
        return 10.0
    return None


def fusion_state(path: Path) -> str:
    text = str(path).lower()
    if "fuse" in text or "融合" in text:
        return "fused_product"
    if "pan" in text and ("mss" in text or "mux" in text):
        return "runtime_fusion_candidate"
    return "unknown"


def iter_image_files(root: Path, limit: int):
    yielded = 0
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames.sort()
        filenames.sort()
        for filename in filenames:
            path = Path(dirpath) / filename
            if path.suffix.lower() not in IMAGE_SUFFIXES:
                continue
            yield path
            yielded += 1
            if yielded >= limit:
                return


def read_preview(path: Path) -> tuple[np.ndarray | None, tuple[int, int, int] | None, str]:
    try:
        import cv2

        arr = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if arr is not None:
            if arr.ndim == 2:
                arr = arr[:, :, None]
            elif arr.ndim == 3 and arr.shape[2] >= 3:
                arr = arr[:, :, :3]
            return arr, normalize_shape(arr), "cv2"
    except Exception:
        pass
    if path.suffix.lower() in {".tif", ".tiff"}:
        try:
            import tifffile

            arr = tifffile.imread(str(path))
            if arr.ndim == 2:
                arr = arr[:, :, None]
            elif arr.ndim == 3 and arr.shape[0] <= 16 and arr.shape[1] > 32 and arr.shape[2] > 32:
                arr = np.moveaxis(arr, 0, -1)
            if arr.ndim == 3:
                arr = arr[:, :, : min(arr.shape[2], 3)]
            return arr, normalize_shape(arr), "tifffile"
        except Exception:
            pass
    try:
        from PIL import Image

        img = Image.open(path)
        arr = np.asarray(img.convert("RGB"))
        return arr, normalize_shape(arr), "pil"
    except Exception:
        return None, None, "unreadable"


def normalize_shape(arr: np.ndarray) -> tuple[int, int, int] | None:
    if arr.ndim != 3:
        return None
    h, w, c = arr.shape
    if h <= 0 or w <= 0 or c <= 0:
        return None
    return int(h), int(w), int(c)


def to_uint8_preview(arr: np.ndarray) -> np.ndarray:
    arr = arr.astype(np.float32)
    out = np.zeros_like(arr, dtype=np.uint8)
    for c in range(arr.shape[2]):
        band = arr[:, :, c]
        lo, hi = np.percentile(band, [2, 98])
        if hi <= lo:
            out[:, :, c] = 0
        else:
            out[:, :, c] = np.clip((band - lo) * 255.0 / (hi - lo), 0, 255).astype(np.uint8)
    return out


def quality(arr: np.ndarray, max_side: int = 256) -> dict[str, float]:
    h, w = arr.shape[:2]
    step_y = max(1, math.ceil(h / max_side))
    step_x = max(1, math.ceil(w / max_side))
    sample = arr[::step_y, ::step_x]
    sample8 = to_uint8_preview(sample)
    gray = sample8.mean(axis=2)
    black = float((gray <= 3).mean())
    white = float((gray >= 252).mean())
    valid = float(((gray > 3) & (gray < 252)).mean())
    std = float(gray.std())
    return {
        "black_ratio": black,
        "white_ratio": white,
        "valid_ratio": valid,
        "std": std,
    }


def make_windows(width: int, height: int, patch: int, stride: int, max_windows: int, rng: random.Random) -> list[dict[str, int]]:
    if width < patch or height < patch:
        return [{"x": 0, "y": 0, "width": width, "height": height}]
    xs = list(range(0, max(width - patch + 1, 1), stride))
    ys = list(range(0, max(height - patch + 1, 1), stride))
    if xs[-1] != width - patch:
        xs.append(width - patch)
    if ys[-1] != height - patch:
        ys.append(height - patch)
    windows = [{"x": x, "y": y, "width": patch, "height": patch} for y in ys for x in xs]
    if len(windows) > max_windows:
        windows = rng.sample(windows, max_windows)
        windows.sort(key=lambda item: (item["y"], item["x"]))
    return windows


def reject(path: Path, reason: str, rows: list[dict[str, Any]], extra: dict[str, Any] | None = None) -> None:
    row = {"path": str(path), "reason": reason}
    if extra:
        row.update(extra)
    rows.append(row)


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 append_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as fp:
        for row in rows:
            fp.write(json.dumps(row, ensure_ascii=False) + "\n")


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    keys = sorted({k for row in rows for k in row})
    with path.open("w", newline="", encoding="utf-8-sig") as fp:
        writer = csv.DictWriter(fp, fieldnames=keys)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    args = parse_args()
    rng = random.Random(args.seed)
    output = Path(args.output_root)
    output.mkdir(parents=True, exist_ok=True)
    stats = Stats()
    accepted_files: list[dict[str, Any]] = []
    rejected: list[dict[str, Any]] = []
    samples: list[dict[str, Any]] = []
    manifest_dir = output / "manifests"
    report_dir = output / "reports"
    sample_path = manifest_dir / "unlabeled_mim_samples.jsonl"
    accepted_path = manifest_dir / "accepted_source_images.jsonl"
    rejected_path = manifest_dir / "rejected_source_images.jsonl"
    for path in (sample_path, accepted_path, rejected_path):
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text("", encoding="utf-8")

    for root_text in args.roots:
        root = Path(root_text)
        if not root.exists():
            continue
        pending_samples: list[dict[str, Any]] = []
        pending_accepted: list[dict[str, Any]] = []
        pending_rejected: list[dict[str, Any]] = []
        for path in iter_image_files(root, args.max_files_per_root):
            if len(samples) >= args.max_total_windows:
                break
            stats.scanned_files += 1
            arr, shape, reader = read_preview(path)
            if arr is None or shape is None:
                stats.unreadable += 1
                stats.rejected_files += 1
                reject(path, "unreadable", pending_rejected)
                continue
            h, w, c = shape
            if min(h, w) < 64:
                stats.too_small += 1
                stats.rejected_files += 1
                reject(path, "too_small", pending_rejected, {"height": h, "width": w})
                continue
            q = quality(arr)
            if q["valid_ratio"] < args.min_valid_ratio:
                stats.too_black += 1
                stats.rejected_files += 1
                reject(path, "low_valid_ratio", pending_rejected, q | {"height": h, "width": w})
                continue
            if q["black_ratio"] > args.max_black_ratio:
                stats.too_black += 1
                stats.rejected_files += 1
                reject(path, "too_black", pending_rejected, q | {"height": h, "width": w})
                continue
            if q["white_ratio"] > args.max_white_ratio:
                stats.too_white += 1
                stats.rejected_files += 1
                reject(path, "too_white", pending_rejected, q | {"height": h, "width": w})
                continue
            if q["std"] < args.min_std:
                stats.low_texture += 1
                stats.rejected_files += 1
                reject(path, "low_texture", pending_rejected, q | {"height": h, "width": w})
                continue

            satellite = infer_satellite(str(path))
            sensor = infer_sensor(str(path))
            file_record = {
                "source_path": str(path),
                "root": str(root),
                "height": h,
                "width": w,
                "channels": c,
                "reader": reader,
                "satellite": satellite,
                "sensor": sensor,
                "resolution_m": infer_resolution_m(satellite, sensor),
                "fusion": {"state": fusion_state(path), "method": "unknown", "persisted": True},
                "quality": q,
            }
            accepted_files.append(file_record)
            pending_accepted.append(file_record)
            stats.accepted_files += 1
            for idx, window in enumerate(make_windows(w, h, args.patch_size, args.stride, args.max_windows_per_image, rng)):
                if len(samples) >= args.max_total_windows:
                    break
                sample_id = f"mim_{len(samples):08d}"
                samples.append(
                    {
                        "sample_id": sample_id,
                        "source_path": str(path),
                        "window": window,
                        "patch_size": args.patch_size,
                        "reader": reader,
                        "satellite": satellite,
                        "sensor": sensor,
                        "resolution_m": infer_resolution_m(satellite, sensor),
                        "fusion": file_record["fusion"],
                        "quality": q,
                        "task_type": "masked_image_modeling",
                        "label_path": None,
                    }
                )
                pending_samples.append(samples[-1])
                stats.accepted_windows += 1
            if stats.scanned_files % args.progress_every == 0:
                append_jsonl(sample_path, pending_samples)
                append_jsonl(accepted_path, pending_accepted)
                append_jsonl(rejected_path, pending_rejected)
                rejected.extend(pending_rejected)
                pending_samples.clear()
                pending_accepted.clear()
                pending_rejected.clear()
                progress = {
                    **asdict(stats),
                    "current_root": str(root),
                    "samples_written": sum(1 for _ in sample_path.open("r", encoding="utf-8")),
                }
                (output / "progress.json").write_text(json.dumps(progress, indent=2, ensure_ascii=False), encoding="utf-8")
                print(json.dumps(progress, ensure_ascii=False), flush=True)
        append_jsonl(sample_path, pending_samples)
        append_jsonl(accepted_path, pending_accepted)
        append_jsonl(rejected_path, pending_rejected)
        rejected.extend(pending_rejected)
    rng.shuffle(samples)
    train_count = int(len(samples) * args.split_train)
    for i, sample in enumerate(samples):
        sample["split"] = "train" if i < train_count else "val"
    samples.sort(key=lambda item: item["sample_id"])

    write_jsonl(sample_path, samples)
    write_jsonl(accepted_path, accepted_files)
    write_jsonl(rejected_path, rejected)
    write_csv(report_dir / "rejected_source_images.csv", rejected)

    summary = {
        **asdict(stats),
        "output_root": str(output),
        "roots": args.roots,
        "patch_size": args.patch_size,
        "stride": args.stride,
        "splits": dict(Counter(sample["split"] for sample in samples)),
        "accepted_by_satellite": dict(Counter(str(row["satellite"]) for row in samples)),
        "accepted_by_sensor": dict(Counter(str(row["sensor"]) for row in samples)),
        "accepted_by_reader": dict(Counter(row["reader"] for row in samples)),
        "quality_policy": {
            "min_valid_ratio": args.min_valid_ratio,
            "max_black_ratio": args.max_black_ratio,
            "max_white_ratio": args.max_white_ratio,
            "min_std": args.min_std,
        },
    }
    (output / "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()