"""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()