405 lines
19 KiB
Python
405 lines
19 KiB
Python
"""High-level GF-3 HH/HV water extraction pipeline."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from argparse import Namespace
|
|
|
|
import numpy as np
|
|
|
|
from .constants import (
|
|
CLASS_HIGH_CONFIDENCE_WATER,
|
|
CLASS_INVALID,
|
|
CLASS_KNOWN_WATER,
|
|
CLASS_LOW_CONFIDENCE_WATER,
|
|
CLASS_NAMES,
|
|
CLASS_NON_WATER,
|
|
CLASS_PADDY_WATER_LIKE,
|
|
)
|
|
from .dltb import DltbConfig, load_dltb_scene_zones
|
|
from .envi import parse_envi, read_dem_for_sar, read_envi_band, write_tif
|
|
from .previews import save_class_preview, save_dltb_zone_preview, save_gray_png, save_mask_png, save_preview
|
|
from .raster_ops import close_mask, fill_small_holes, open_mask, otsu_threshold, remove_small_components, robust_normalize, slope_degrees, to_db
|
|
from .vector_io import rasterize_geometries, rasterize_vector_mask, rasterize_water_prior
|
|
from .vector_products import write_classified_vectors
|
|
|
|
|
|
def _load_dltb_masks(args: Namespace, info) -> tuple[dict | None, np.ndarray, np.ndarray, np.ndarray]:
|
|
water_mask = np.zeros((info.lines, info.samples), dtype=bool)
|
|
paddy_mask = np.zeros((info.lines, info.samples), dtype=bool)
|
|
strict_mask = np.zeros((info.lines, info.samples), dtype=bool)
|
|
|
|
if args.dltb_cache_dir is not None and args.dltb_mode != "off":
|
|
water_cache = args.dltb_cache_dir / "water_prior.shp"
|
|
paddy_cache = args.dltb_cache_dir / "paddy.shp"
|
|
strict_cache = args.dltb_cache_dir / "strict_review.shp"
|
|
if water_cache.exists():
|
|
water_mask = rasterize_vector_mask([water_cache], info, 0.0, "--dltb-cache-dir water_prior")
|
|
if paddy_cache.exists():
|
|
paddy_mask = rasterize_vector_mask([paddy_cache], info, 0.0, "--dltb-cache-dir paddy")
|
|
if strict_cache.exists():
|
|
strict_mask = rasterize_vector_mask([strict_cache], info, 0.0, "--dltb-cache-dir strict_review")
|
|
stats = {
|
|
"cache_dir": str(args.dltb_cache_dir),
|
|
"mode": args.dltb_mode,
|
|
"source": "cache",
|
|
"water_prior_path": str(water_cache) if water_cache.exists() else None,
|
|
"paddy_path": str(paddy_cache) if paddy_cache.exists() else None,
|
|
"strict_review_path": str(strict_cache) if strict_cache.exists() else None,
|
|
}
|
|
return stats, water_mask, paddy_mask, strict_mask
|
|
|
|
if args.dltb_gdb is not None and args.dltb_mode != "off":
|
|
zones = load_dltb_scene_zones(
|
|
DltbConfig(gdb=args.dltb_gdb, layer=args.dltb_layer, field=args.dltb_field, mode=args.dltb_mode, max_features=args.dltb_max_features),
|
|
info.bounds,
|
|
)
|
|
water_mask = rasterize_geometries(zones.water_geoms, info)
|
|
paddy_mask = rasterize_geometries(zones.paddy_geoms, info)
|
|
strict_mask = rasterize_geometries(zones.strict_geoms, info)
|
|
stats = {
|
|
"gdb": str(args.dltb_gdb),
|
|
"source": "gdb",
|
|
"layer": args.dltb_layer,
|
|
"field": args.dltb_field,
|
|
"mode": args.dltb_mode,
|
|
"source_crs": zones.crs,
|
|
"features_in_scene": zones.feature_count,
|
|
"class_counts": zones.class_counts,
|
|
"dlmc_values_in_scene": zones.dlmc_values,
|
|
}
|
|
return stats, water_mask, paddy_mask, strict_mask
|
|
|
|
return None, water_mask, paddy_mask, strict_mask
|
|
|
|
|
|
def _ensure_matching_grids(hh_info, hv_info) -> None:
|
|
hh_grid = (hh_info.samples, hh_info.lines, hh_info.x0, hh_info.y0, hh_info.dx, hh_info.dy)
|
|
hv_grid = (hv_info.samples, hv_info.lines, hv_info.x0, hv_info.y0, hv_info.dx, hv_info.dy)
|
|
if hh_grid != hv_grid:
|
|
raise ValueError("HH and HV grids do not match")
|
|
|
|
|
|
def _classify_products(mask: np.ndarray, low_confidence_candidate: np.ndarray, paddy_candidate: np.ndarray, prior_mask: np.ndarray, valid: np.ndarray) -> np.ndarray:
|
|
high_confidence_mask = mask & ~paddy_candidate
|
|
known_water_mask = mask & prior_mask
|
|
classified = np.full(mask.shape, CLASS_INVALID, dtype=np.uint8)
|
|
classified[valid] = CLASS_NON_WATER
|
|
classified[low_confidence_candidate & valid & ~mask & ~paddy_candidate] = CLASS_LOW_CONFIDENCE_WATER
|
|
classified[high_confidence_mask & valid] = CLASS_HIGH_CONFIDENCE_WATER
|
|
classified[paddy_candidate & valid] = CLASS_PADDY_WATER_LIKE
|
|
classified[known_water_mask & valid] = CLASS_KNOWN_WATER
|
|
return classified
|
|
|
|
|
|
def _write_rasters(
|
|
args: Namespace,
|
|
info,
|
|
probability: np.ndarray,
|
|
valid: np.ndarray,
|
|
mask: np.ndarray,
|
|
raw_mask: np.ndarray,
|
|
classified: np.ndarray,
|
|
cartographic_water: np.ndarray | None,
|
|
prior_mask: np.ndarray,
|
|
paddy_mask: np.ndarray,
|
|
paddy_candidate: np.ndarray,
|
|
dltb_enabled: bool,
|
|
dltb_water_mask: np.ndarray,
|
|
dltb_paddy_mask: np.ndarray,
|
|
dltb_strict_mask: np.ndarray,
|
|
) -> None:
|
|
mask_u8 = np.where(valid, mask.astype(np.uint8), 255).astype(np.uint8)
|
|
raw_mask_u8 = np.where(valid, raw_mask.astype(np.uint8), 255).astype(np.uint8)
|
|
|
|
write_tif(args.out_dir / "water_score.tif", np.where(np.isfinite(probability), probability, -9999.0), info, "float32", nodata=-9999.0)
|
|
write_tif(args.out_dir / "water_mask.tif", mask_u8, info, "uint8", nodata=255)
|
|
write_tif(args.out_dir / "water_mask_raw.tif", raw_mask_u8, info, "uint8", nodata=255)
|
|
write_tif(args.out_dir / "classified_water.tif", classified, info, "uint8", nodata=CLASS_INVALID)
|
|
if cartographic_water is not None:
|
|
write_tif(args.out_dir / "cartographic_water.tif", np.where(valid, cartographic_water.astype(np.uint8), 255).astype(np.uint8), info, "uint8", nodata=255)
|
|
if args.water_vector:
|
|
write_tif(args.out_dir / "known_water_prior.tif", prior_mask.astype(np.uint8), info, "uint8", nodata=0)
|
|
if args.paddy_vector or dltb_enabled:
|
|
write_tif(args.out_dir / "paddy_prior.tif", paddy_mask.astype(np.uint8), info, "uint8", nodata=0)
|
|
write_tif(args.out_dir / "paddy_water_like.tif", np.where(valid, paddy_candidate.astype(np.uint8), 255).astype(np.uint8), info, "uint8", nodata=255)
|
|
if dltb_enabled:
|
|
write_tif(args.out_dir / "dltb_water_prior.tif", dltb_water_mask.astype(np.uint8), info, "uint8", nodata=0)
|
|
write_tif(args.out_dir / "dltb_paddy_prior.tif", dltb_paddy_mask.astype(np.uint8), info, "uint8", nodata=0)
|
|
write_tif(args.out_dir / "dltb_strict_zone.tif", dltb_strict_mask.astype(np.uint8), info, "uint8", nodata=0)
|
|
|
|
|
|
def _write_previews(
|
|
args: Namespace,
|
|
probability: np.ndarray,
|
|
valid: np.ndarray,
|
|
mask: np.ndarray,
|
|
raw_mask: np.ndarray,
|
|
hh_norm: np.ndarray,
|
|
hv_norm: np.ndarray,
|
|
classified: np.ndarray,
|
|
prior_mask: np.ndarray,
|
|
paddy_candidate: np.ndarray,
|
|
dltb_enabled: bool,
|
|
dltb_water_mask: np.ndarray,
|
|
dltb_paddy_mask: np.ndarray,
|
|
dltb_strict_mask: np.ndarray,
|
|
cartographic_water: np.ndarray | None,
|
|
) -> None:
|
|
save_gray_png(args.out_dir / "water_score.png", probability, valid)
|
|
save_mask_png(args.out_dir / "water_mask.png", mask)
|
|
save_mask_png(args.out_dir / "water_mask_raw.png", raw_mask)
|
|
if args.water_vector:
|
|
save_mask_png(args.out_dir / "known_water_prior.png", prior_mask)
|
|
if args.paddy_vector or dltb_enabled:
|
|
save_mask_png(args.out_dir / "paddy_water_like.png", paddy_candidate)
|
|
if dltb_enabled:
|
|
save_dltb_zone_preview(args.out_dir / "dltb_zone_preview.png", dltb_water_mask, dltb_paddy_mask, dltb_strict_mask)
|
|
if cartographic_water is not None:
|
|
save_mask_png(args.out_dir / "cartographic_water.png", cartographic_water)
|
|
save_preview(args.out_dir / "preview_overlay.png", hh_norm, hv_norm, mask, valid)
|
|
save_class_preview(args.out_dir / "classified_preview.png", classified, hh_norm, hv_norm, valid)
|
|
|
|
|
|
def run_from_args(args: Namespace) -> int:
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
hh_info = parse_envi(args.hh)
|
|
hv_info = parse_envi(args.hv)
|
|
_ensure_matching_grids(hh_info, hv_info)
|
|
|
|
dltb_stats, dltb_water_mask, dltb_paddy_mask, dltb_strict_mask = _load_dltb_masks(args, hh_info)
|
|
dltb_enabled = dltb_stats is not None
|
|
|
|
hh = read_envi_band(hh_info)
|
|
hv = read_envi_band(hv_info)
|
|
valid = np.isfinite(hh) & np.isfinite(hv) & (hh > 0) & (hv > 0)
|
|
|
|
hh_db = to_db(hh)
|
|
hv_db = to_db(hv)
|
|
valid &= np.isfinite(hh_db) & np.isfinite(hv_db)
|
|
hh_norm, hh_lo, hh_hi = robust_normalize(hh_db, valid)
|
|
hv_norm, hv_lo, hv_hi = robust_normalize(hv_db, valid)
|
|
|
|
low_backscatter = 1.0 - 0.5 * (hh_norm + hv_norm)
|
|
low_backscatter[~valid] = np.nan
|
|
|
|
values = low_backscatter[valid]
|
|
if args.threshold_method == "otsu":
|
|
score_threshold = otsu_threshold(values)
|
|
else:
|
|
score_threshold = float(np.nanpercentile(values, args.score_percentile))
|
|
hv_dark_threshold = float(np.nanpercentile(hv_norm[valid], args.hv_percentile))
|
|
|
|
mask = (low_backscatter >= score_threshold) & (hv_norm <= hv_dark_threshold) & valid
|
|
prior_mask = np.zeros(mask.shape, dtype=bool)
|
|
prior_candidate_pixels = 0
|
|
prior_score_threshold = None
|
|
prior_hv_threshold = None
|
|
if args.water_vector:
|
|
prior_mask = rasterize_water_prior(args.water_vector, hh_info, args.river_buffer_meters) & valid
|
|
prior_score_threshold = float(np.nanpercentile(values, args.prior_score_percentile))
|
|
prior_hv_threshold = float(np.nanpercentile(hv_norm[valid], args.prior_hv_percentile))
|
|
prior_candidate = prior_mask & (low_backscatter >= prior_score_threshold) & (hv_norm <= prior_hv_threshold)
|
|
prior_candidate_pixels = int(prior_candidate.sum())
|
|
mask |= prior_candidate
|
|
if dltb_enabled:
|
|
prior_mask |= dltb_water_mask & valid
|
|
prior_score_threshold = prior_score_threshold if prior_score_threshold is not None else float(np.nanpercentile(values, args.prior_score_percentile))
|
|
prior_hv_threshold = prior_hv_threshold if prior_hv_threshold is not None else float(np.nanpercentile(hv_norm[valid], args.prior_hv_percentile))
|
|
dltb_prior_candidate = prior_mask & (low_backscatter >= prior_score_threshold) & (hv_norm <= prior_hv_threshold)
|
|
prior_candidate_pixels += int((dltb_water_mask & dltb_prior_candidate).sum())
|
|
mask |= dltb_prior_candidate
|
|
|
|
paddy_mask = np.zeros(mask.shape, dtype=bool)
|
|
paddy_candidate = np.zeros(mask.shape, dtype=bool)
|
|
paddy_score_threshold = None
|
|
paddy_hv_threshold = None
|
|
if args.paddy_vector:
|
|
paddy_mask = rasterize_vector_mask(args.paddy_vector, hh_info, args.paddy_buffer_meters, "--paddy-vector") & valid
|
|
paddy_score_threshold = float(np.nanpercentile(values, args.paddy_score_percentile))
|
|
paddy_hv_threshold = float(np.nanpercentile(hv_norm[valid], args.paddy_hv_percentile))
|
|
paddy_candidate = paddy_mask & (low_backscatter >= paddy_score_threshold) & (hv_norm <= paddy_hv_threshold)
|
|
if dltb_enabled:
|
|
paddy_mask |= dltb_paddy_mask & valid
|
|
paddy_score_threshold = paddy_score_threshold if paddy_score_threshold is not None else float(np.nanpercentile(values, args.paddy_score_percentile))
|
|
paddy_hv_threshold = paddy_hv_threshold if paddy_hv_threshold is not None else float(np.nanpercentile(hv_norm[valid], args.paddy_hv_percentile))
|
|
paddy_candidate |= paddy_mask & (low_backscatter >= paddy_score_threshold) & (hv_norm <= paddy_hv_threshold)
|
|
|
|
candidate_score_threshold = float(np.nanpercentile(values, args.candidate_score_percentile))
|
|
candidate_hv_threshold = float(np.nanpercentile(hv_norm[valid], args.candidate_hv_percentile))
|
|
low_confidence_candidate = (low_backscatter >= candidate_score_threshold) & (hv_norm <= candidate_hv_threshold) & valid
|
|
if dltb_enabled and args.dltb_mode == "soft":
|
|
strong_candidate = (low_backscatter >= score_threshold) & (hv_norm <= hv_dark_threshold) & valid
|
|
low_confidence_candidate &= ~dltb_strict_mask | strong_candidate
|
|
elif dltb_enabled and args.dltb_mode == "strict":
|
|
low_confidence_candidate &= ~dltb_strict_mask
|
|
|
|
dem_used = False
|
|
slope_threshold = None
|
|
slope = None
|
|
if args.dem is not None:
|
|
dem_info = parse_envi(args.dem)
|
|
dem = read_dem_for_sar(dem_info, hh_info)
|
|
slope = slope_degrees(dem, hh_info.dx, hh_info.dy, center_lat=hh_info.y0 - hh_info.lines * hh_info.dy * 0.5)
|
|
slope_threshold = float(args.slope_max)
|
|
mask &= np.isfinite(slope) & (slope <= args.slope_max)
|
|
paddy_candidate &= np.isfinite(slope) & (slope <= args.slope_max)
|
|
low_confidence_candidate &= np.isfinite(slope) & (slope <= args.slope_max)
|
|
valid &= np.isfinite(slope)
|
|
dem_used = True
|
|
save_gray_png(args.out_dir / "slope_preview.png", slope, np.isfinite(slope))
|
|
write_tif(args.out_dir / "slope_degrees.tif", np.where(np.isfinite(slope), slope, -9999.0), hh_info, "float32", nodata=-9999.0)
|
|
|
|
raw_mask = mask.copy()
|
|
if not args.no_morphology:
|
|
mask = close_mask(mask, args.close_pixels, valid)
|
|
mask = remove_small_components(mask, args.min_component_pixels)
|
|
mask = fill_small_holes(mask, args.fill_hole_pixels)
|
|
mask = open_mask(mask, args.open_pixels, valid)
|
|
paddy_candidate = close_mask(paddy_candidate, args.paddy_close_pixels, valid)
|
|
paddy_candidate = open_mask(paddy_candidate, args.paddy_open_pixels, valid)
|
|
low_confidence_candidate = close_mask(low_confidence_candidate, args.candidate_close_pixels, valid)
|
|
low_confidence_candidate = open_mask(low_confidence_candidate, args.candidate_open_pixels, valid)
|
|
|
|
probability = np.clip(low_backscatter, 0.0, 1.0)
|
|
probability[~valid] = np.nan
|
|
classified = _classify_products(mask, low_confidence_candidate, paddy_candidate, prior_mask, valid)
|
|
|
|
cartographic_water = None
|
|
if args.cartographic_water:
|
|
cartographic_water = (mask | low_confidence_candidate) & valid
|
|
if args.cartographic_include_paddy:
|
|
cartographic_water |= paddy_candidate & valid
|
|
else:
|
|
cartographic_water &= ~paddy_candidate
|
|
if not args.no_morphology:
|
|
cartographic_water = close_mask(cartographic_water, args.cartographic_close_pixels, valid)
|
|
cartographic_water = fill_small_holes(cartographic_water, args.cartographic_fill_hole_pixels)
|
|
cartographic_water = remove_small_components(cartographic_water, args.cartographic_min_component_pixels)
|
|
cartographic_water = open_mask(cartographic_water, args.cartographic_open_pixels, valid)
|
|
|
|
_write_rasters(
|
|
args,
|
|
hh_info,
|
|
probability,
|
|
valid,
|
|
mask,
|
|
raw_mask,
|
|
classified,
|
|
cartographic_water,
|
|
prior_mask,
|
|
paddy_mask,
|
|
paddy_candidate,
|
|
dltb_enabled,
|
|
dltb_water_mask,
|
|
dltb_paddy_mask,
|
|
dltb_strict_mask,
|
|
)
|
|
_write_previews(
|
|
args,
|
|
probability,
|
|
valid,
|
|
mask,
|
|
raw_mask,
|
|
hh_norm,
|
|
hv_norm,
|
|
classified,
|
|
prior_mask,
|
|
paddy_candidate,
|
|
dltb_enabled,
|
|
dltb_water_mask,
|
|
dltb_paddy_mask,
|
|
dltb_strict_mask,
|
|
cartographic_water,
|
|
)
|
|
|
|
vector_stats = None
|
|
if args.out_vector_gpkg is not None or args.out_vector_shp_dir is not None:
|
|
vector_stats = write_classified_vectors(
|
|
args.out_vector_gpkg,
|
|
args.out_vector_shp_dir,
|
|
classified,
|
|
cartographic_water,
|
|
hh_info,
|
|
probability,
|
|
hh_db,
|
|
hv_db,
|
|
slope,
|
|
prior_mask,
|
|
paddy_mask,
|
|
args.min_polygon_area_m2,
|
|
args.simplify_meters,
|
|
args.smooth_meters,
|
|
args.min_hole_area_m2,
|
|
)
|
|
|
|
valid_count = int(valid.sum())
|
|
stats = {
|
|
"hh": str(args.hh),
|
|
"hv": str(args.hv),
|
|
"dem": str(args.dem) if args.dem else None,
|
|
"dltb": dltb_stats,
|
|
"water_vectors": [str(p) for p in args.water_vector],
|
|
"paddy_vectors": [str(p) for p in args.paddy_vector],
|
|
"river_buffer_meters": float(args.river_buffer_meters),
|
|
"paddy_buffer_meters": float(args.paddy_buffer_meters),
|
|
"shape": [hh_info.lines, hh_info.samples],
|
|
"bounds_wgs84": list(hh_info.bounds),
|
|
"valid_pixels": valid_count,
|
|
"valid_ratio": float(valid_count / valid.size),
|
|
"hh_db_percentile_2_98": [hh_lo, hh_hi],
|
|
"hv_db_percentile_2_98": [hv_lo, hv_hi],
|
|
"threshold_method": args.threshold_method,
|
|
"score_threshold": float(score_threshold),
|
|
"hv_norm_dark_threshold": float(hv_dark_threshold),
|
|
"prior_score_threshold": prior_score_threshold,
|
|
"prior_hv_norm_dark_threshold": prior_hv_threshold,
|
|
"paddy_score_threshold": paddy_score_threshold,
|
|
"paddy_hv_norm_dark_threshold": paddy_hv_threshold,
|
|
"candidate_score_threshold": candidate_score_threshold,
|
|
"candidate_hv_norm_dark_threshold": candidate_hv_threshold,
|
|
"slope_threshold_degrees": slope_threshold,
|
|
"dem_used": dem_used,
|
|
"known_water_prior_pixels": int(prior_mask.sum()),
|
|
"known_water_prior_ratio_valid": float(prior_mask.sum() / max(valid_count, 1)),
|
|
"dltb_water_prior_pixels": int(dltb_water_mask.sum()),
|
|
"dltb_paddy_prior_pixels": int(dltb_paddy_mask.sum()),
|
|
"dltb_strict_zone_pixels": int(dltb_strict_mask.sum()),
|
|
"prior_candidate_pixels": prior_candidate_pixels,
|
|
"paddy_prior_pixels": int(paddy_mask.sum()),
|
|
"paddy_water_like_pixels": int(paddy_candidate.sum()),
|
|
"paddy_water_like_ratio_valid": float(paddy_candidate.sum() / max(valid_count, 1)),
|
|
"low_confidence_water_pixels": int((classified == CLASS_LOW_CONFIDENCE_WATER).sum()),
|
|
"cartographic_water_pixels": int(cartographic_water.sum()) if cartographic_water is not None else 0,
|
|
"cartographic_water_ratio_valid": float(cartographic_water.sum() / max(valid_count, 1)) if cartographic_water is not None else 0.0,
|
|
"raw_water_pixels": int(raw_mask.sum()),
|
|
"raw_water_ratio_valid": float(raw_mask.sum() / max(valid_count, 1)),
|
|
"water_pixels": int(mask.sum()),
|
|
"water_ratio_valid": float(mask.sum() / max(valid_count, 1)),
|
|
"classified_counts": {CLASS_NAMES[class_id]: int((classified == class_id).sum()) for class_id in CLASS_NAMES},
|
|
"min_component_pixels": int(args.min_component_pixels),
|
|
"fill_hole_pixels": int(args.fill_hole_pixels),
|
|
"close_pixels": int(args.close_pixels),
|
|
"open_pixels": int(args.open_pixels),
|
|
"paddy_close_pixels": int(args.paddy_close_pixels),
|
|
"paddy_open_pixels": int(args.paddy_open_pixels),
|
|
"candidate_close_pixels": int(args.candidate_close_pixels),
|
|
"candidate_open_pixels": int(args.candidate_open_pixels),
|
|
"cartographic_water_enabled": bool(args.cartographic_water),
|
|
"cartographic_include_paddy": bool(args.cartographic_include_paddy),
|
|
"cartographic_close_pixels": int(args.cartographic_close_pixels),
|
|
"cartographic_open_pixels": int(args.cartographic_open_pixels),
|
|
"cartographic_fill_hole_pixels": int(args.cartographic_fill_hole_pixels),
|
|
"cartographic_min_component_pixels": int(args.cartographic_min_component_pixels),
|
|
"min_polygon_area_m2": float(args.min_polygon_area_m2),
|
|
"simplify_meters": float(args.simplify_meters),
|
|
"smooth_meters": float(args.smooth_meters),
|
|
"min_hole_area_m2": float(args.min_hole_area_m2),
|
|
"vector_output": vector_stats,
|
|
}
|
|
(args.out_dir / "metadata.json").write_text(json.dumps(stats, indent=2), encoding="utf-8")
|
|
print(json.dumps(stats, indent=2))
|
|
return 0
|
|
|