"""Run a small, inspectable CPU semantic-mapping baseline.
|
|
The first capability version deliberately uses deterministic RGB/HSV rules so it
|
can run without downloading model weights. GeoAI's ``masks_to_vector`` is used
|
for class-wise polygon extraction, while the metadata keeps the model boundary
|
explicit: this is a B capability, not a built-in GeoAI segmentation model.
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
import json
|
import time
|
from datetime import UTC, datetime
|
from pathlib import Path
|
from typing import Any
|
|
import cv2
|
import geopandas as gpd
|
import numpy as np
|
import pandas as pd
|
import rasterio
|
from PIL import Image
|
from rasterio.features import shapes
|
from rasterio.transform import Affine
|
from shapely.geometry import shape
|
|
|
CLASS_INFO = {
|
0: {"key": "other", "label": "其他", "color": (80, 80, 80)},
|
1: {"key": "vegetation", "label": "植被", "color": (40, 180, 70)},
|
2: {"key": "water", "label": "水体", "color": (35, 130, 220)},
|
3: {"key": "impervious", "label": "不透水面", "color": (220, 150, 55)},
|
}
|
SUPPORTED_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
|
|
|
def classify_rgb(rgb: np.ndarray) -> np.ndarray:
|
"""Classify RGB pixels with conservative, explainable colour rules."""
|
rgb_f = rgb.astype(np.float32)
|
r, g, b = [rgb_f[..., i] for i in range(3)]
|
hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
|
saturation = hsv[..., 1].astype(np.float32)
|
value = hsv[..., 2].astype(np.float32)
|
result = np.zeros(rgb.shape[:2], dtype=np.uint8)
|
|
green = (g > r * 1.04) & (g > b * 1.04) & (g > 45) & (saturation > 45)
|
blue = (b > r * 1.08) & (b > g * 1.02) & (b > 45) & (saturation > 35)
|
gray = (np.max(rgb_f, axis=2) - np.min(rgb_f, axis=2) < 38) & (value > 45)
|
result[green] = 1
|
result[blue & ~green] = 2
|
result[gray & ~green & ~blue] = 3
|
return result
|
|
|
def _raster_profile(input_path: Path, width: int, height: int) -> tuple[Affine, str | None, bool]:
|
if input_path.suffix.lower() not in {".tif", ".tiff"}:
|
return Affine.identity(), None, False
|
try:
|
with rasterio.open(input_path) as src:
|
transform = src.transform
|
crs = src.crs.to_string() if src.crs else None
|
return transform, crs, bool(src.crs and src.transform != Affine.identity())
|
except rasterio.errors.RasterioIOError:
|
return Affine.identity(), None, False
|
|
|
def _geojson_from_mask(mask: np.ndarray, transform: Affine, crs: str | None, class_id: int) -> gpd.GeoDataFrame:
|
records: list[dict[str, Any]] = []
|
binary = (mask == class_id).astype(np.uint8)
|
for geometry, value in shapes(binary, mask=binary.astype(bool), transform=transform):
|
if value != 1:
|
continue
|
polygon = shape(geometry)
|
if polygon.area < 100:
|
continue
|
records.append({"geometry": polygon, "class_id": class_id, "class_key": CLASS_INFO[class_id]["key"], "class_label": CLASS_INFO[class_id]["label"], "pixel_area": float(polygon.area)})
|
return gpd.GeoDataFrame(records, geometry="geometry", crs=crs)
|
|
|
def vectorize_with_geoai(mask: np.ndarray, transform: Affine, crs: str | None, output_path: Path) -> bool:
|
"""Use GeoAI's class-wise mask utility, then preserve class properties."""
|
try:
|
from geoai import masks_to_vector
|
|
frames: list[gpd.GeoDataFrame] = []
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
for class_id in (1, 2, 3):
|
temp_path = output_path.with_name(f".{output_path.stem}-{class_id}.tif")
|
binary = np.where(mask == class_id, 255, 0).astype(np.uint8)
|
with rasterio.open(temp_path, "w", driver="GTiff", height=binary.shape[0], width=binary.shape[1], count=1, dtype="uint8", transform=transform, crs=crs, nodata=0) as dst:
|
dst.write(binary, 1)
|
try:
|
frame = masks_to_vector(str(temp_path), min_object_area=100, simplify_tolerance=1.0)
|
if not frame.empty:
|
frame = frame.assign(class_id=class_id, class_key=CLASS_INFO[class_id]["key"], class_label=CLASS_INFO[class_id]["label"])
|
frames.append(frame[["geometry", "class_id", "class_key", "class_label", "confidence"]])
|
finally:
|
temp_path.unlink(missing_ok=True)
|
merged = gpd.GeoDataFrame(pd_concat(frames), geometry="geometry", crs=crs)
|
merged.to_file(output_path, driver="GeoJSON")
|
return True
|
except Exception:
|
fallback = gpd.GeoDataFrame(pd_concat([_geojson_from_mask(mask, transform, crs, class_id) for class_id in (1, 2, 3)]), geometry="geometry", crs=crs)
|
fallback.to_file(output_path, driver="GeoJSON")
|
return False
|
|
|
def pd_concat(frames: list[gpd.GeoDataFrame]) -> gpd.GeoDataFrame:
|
if not frames:
|
return gpd.GeoDataFrame({"geometry": []}, geometry="geometry")
|
return gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), geometry="geometry", crs=frames[0].crs)
|
|
|
def process_image(input_path: Path, output_dir: Path) -> dict[str, Any]:
|
started = time.perf_counter()
|
with Image.open(input_path) as image:
|
rgb = np.asarray(image.convert("RGB"))
|
height, width = rgb.shape[:2]
|
mask = classify_rgb(rgb)
|
transform, crs, georeferenced = _raster_profile(input_path, width, height)
|
stem = input_path.stem
|
mask_path = output_dir / f"{stem}.mask.png"
|
overlay_path = output_dir / f"{stem}.overlay.png"
|
raster_path = output_dir / f"{stem}.mask.tif"
|
vector_path = output_dir / f"{stem}.segments.geojson"
|
Image.fromarray(mask).save(mask_path)
|
palette = np.zeros((256, 3), dtype=np.uint8)
|
for class_id, info in CLASS_INFO.items():
|
palette[class_id] = info["color"]
|
color_mask = palette[mask]
|
overlay = (rgb.astype(np.float32) * 0.52 + color_mask.astype(np.float32) * 0.48).clip(0, 255).astype(np.uint8)
|
Image.fromarray(overlay).save(overlay_path, quality=92)
|
with rasterio.open(raster_path, "w", driver="GTiff", height=height, width=width, count=1, dtype="uint8", transform=transform, crs=crs, nodata=0) as dst:
|
dst.write(mask, 1)
|
used_geoai = vectorize_with_geoai(mask, transform, crs, vector_path)
|
counts = {info["key"]: int(np.count_nonzero(mask == class_id)) for class_id, info in CLASS_INFO.items()}
|
return {"file": input_path.name, "width": width, "height": height, "mask_file": mask_path.name, "overlay_file": overlay_path.name, "raster_file": raster_path.name, "vector_file": vector_path.name, "class_pixel_counts": counts, "georeferenced": georeferenced, "vectorizer": "geoai.masks_to_vector" if used_geoai else "rasterio.features.shapes fallback", "elapsed_seconds": round(time.perf_counter() - started, 3)}
|
|
|
def collect_inputs(input_path: Path) -> list[Path]:
|
if input_path.is_file():
|
candidates = [input_path]
|
elif input_path.is_dir():
|
candidates = sorted(item for item in input_path.iterdir() if item.is_file())
|
else:
|
raise SystemExit(f"Input path does not exist: {input_path}")
|
inputs = [item for item in candidates if item.suffix.lower() in SUPPORTED_SUFFIXES]
|
if not inputs:
|
raise SystemExit("No JPG, PNG, or GeoTIFF inputs were found.")
|
return inputs
|
|
|
def main() -> int:
|
parser = argparse.ArgumentParser(description="Run the CPU semantic-mapping baseline.")
|
parser.add_argument("--input", type=Path, default=Path("shared/data/raw/01-object-detection"))
|
parser.add_argument("--output", type=Path, default=Path("shared/outputs/02-semantic-mapping"))
|
args = parser.parse_args()
|
inputs = collect_inputs(args.input)
|
if args.output.exists() and any(args.output.iterdir()):
|
raise SystemExit(f"Output directory is not empty: {args.output}. Use a new run directory.")
|
args.output.mkdir(parents=True, exist_ok=True)
|
started = time.perf_counter()
|
images = [process_image(item, args.output) for item in inputs]
|
metadata = {"capability": "02-semantic-mapping", "classification": "B", "created_at": datetime.now(UTC).isoformat(), "geoai_version": __import__("importlib.metadata").metadata.version("geoai-py"), "method": "deterministic RGB/HSV semantic baseline", "model": "classical-color-baseline", "device": "CPU", "thresholds": {"vegetation_green_ratio": 1.04, "water_blue_ratio": 1.08, "gray_channel_range": 38, "minimum_vector_area_pixels": 100}, "input_dir": args.input.as_posix(), "raw_input_dir": args.input.as_posix(), "input_count": len(images), "processed_images": len(images), "elapsed_seconds": round(time.perf_counter() - started, 3), "images": images, "classes": [info | {"id": class_id} for class_id, info in CLASS_INFO.items()], "limitations": ["这是可解释的颜色规则基线，不是经过训练的通用分割模型。", "普通 JPG/PNG 没有 CRS 时，GeoJSON 使用像素坐标；只有有效 GeoTIFF 地理参考才会保留地理坐标。", "复杂光照、阴影、材质相近区域可能误分；当前没有人工标注，因此不报告精度。"]}
|
(args.output / "run_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
|
print(json.dumps(metadata, ensure_ascii=False, indent=2))
|
return 0
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
