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