Align erosion to mask with order-only rivers
This commit is contained in:
@@ -187,7 +187,7 @@ def _layer_has_field(layer: ogr.Layer, field_name: Optional[str]) -> bool:
|
||||
return layer_defn.GetFieldIndex(field_name) != -1
|
||||
|
||||
|
||||
def _build_depth_lines(
|
||||
def _build_order_lines(
|
||||
layer_sources: List[Tuple[ogr.Layer, str]],
|
||||
bounds: Tuple[float, float, float, float],
|
||||
width: int,
|
||||
@@ -199,10 +199,10 @@ def _build_depth_lines(
|
||||
) -> np.ndarray:
|
||||
gt = _tile_geotransform(bounds, width, height)
|
||||
driver = gdal.GetDriverByName("MEM")
|
||||
depth_ds = driver.Create("", width, height, 1, gdal.GDT_Float32)
|
||||
depth_ds.SetGeoTransform(gt)
|
||||
depth_ds.SetProjection(tile_srs.ExportToWkt())
|
||||
band = depth_ds.GetRasterBand(1)
|
||||
order_ds = driver.Create("", width, height, 1, gdal.GDT_Float32)
|
||||
order_ds.SetGeoTransform(gt)
|
||||
order_ds.SetProjection(tile_srs.ExportToWkt())
|
||||
band = order_ds.GetRasterBand(1)
|
||||
band.Fill(0.0)
|
||||
|
||||
if not any(_layer_has_field(layer, profile.order_field) for layer, _ in layer_sources):
|
||||
@@ -222,9 +222,6 @@ def _build_depth_lines(
|
||||
order = _normalize_order(feat.GetField(profile.order_field), profile)
|
||||
if order is None:
|
||||
continue
|
||||
depth_m = _order_depth(order, profile)
|
||||
if depth_m is None:
|
||||
continue
|
||||
geom = feat.GetGeometryRef()
|
||||
if geom is None:
|
||||
continue
|
||||
@@ -233,7 +230,7 @@ def _build_depth_lines(
|
||||
if geom.IsEmpty():
|
||||
continue
|
||||
buffered_geoms.append(geom)
|
||||
burn_values.append(float(depth_m))
|
||||
burn_values.append(float(order))
|
||||
|
||||
if not buffered_geoms:
|
||||
return band.ReadAsArray().astype(np.float32)
|
||||
@@ -244,7 +241,7 @@ def _build_depth_lines(
|
||||
|
||||
if hasattr(gdal, "RasterizeGeometries"):
|
||||
gdal.RasterizeGeometries(
|
||||
depth_ds,
|
||||
order_ds,
|
||||
[1],
|
||||
ordered_geoms,
|
||||
burn_values=ordered_values,
|
||||
@@ -265,20 +262,14 @@ def _build_depth_lines(
|
||||
mem_layer.CreateFeature(feature)
|
||||
feature = None
|
||||
gdal.RasterizeLayer(
|
||||
depth_ds,
|
||||
order_ds,
|
||||
[1],
|
||||
mem_layer,
|
||||
options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"],
|
||||
)
|
||||
mem_ds = None
|
||||
|
||||
depth = band.ReadAsArray().astype(np.float32)
|
||||
if profile.smooth_sigma_m > 0.0:
|
||||
sigma_px = profile.smooth_sigma_m / max(_pixel_size_m(gt), 1e-6)
|
||||
if sigma_px > 0.0:
|
||||
depth = ndimage.gaussian_filter(depth, sigma=sigma_px, mode="nearest")
|
||||
|
||||
return depth
|
||||
return band.ReadAsArray().astype(np.float32)
|
||||
|
||||
|
||||
def _build_flow_distance(
|
||||
@@ -728,7 +719,7 @@ def erode_rivers(cfg: Config) -> int:
|
||||
else:
|
||||
mask = _refine_mask(mask, fill_radius, smooth_sigma)
|
||||
|
||||
depth_line = _build_depth_lines(
|
||||
order_line = _build_order_lines(
|
||||
layer_sources,
|
||||
mosaic_bounds,
|
||||
mosaic_w,
|
||||
@@ -739,42 +730,47 @@ def erode_rivers(cfg: Config) -> int:
|
||||
tile_srs,
|
||||
)
|
||||
|
||||
line_mask = depth_line > 0.0
|
||||
line_mask = order_line > 0.0
|
||||
mask_bin = mask > 0.1
|
||||
depth_map = np.zeros_like(depth_line, dtype=np.float32)
|
||||
depth_map = np.zeros_like(order_line, dtype=np.float32)
|
||||
labels = None
|
||||
num = 0
|
||||
indices = None
|
||||
depth_nearest = None
|
||||
order_nearest = None
|
||||
dist_nearest = None
|
||||
|
||||
if np.any(line_mask):
|
||||
_, indices = ndimage.distance_transform_edt(~line_mask, return_indices=True)
|
||||
order_nearest = order_line[tuple(indices)]
|
||||
|
||||
if np.any(mask_bin):
|
||||
labels, num = ndimage.label(mask_bin)
|
||||
fallback_depth = float(getattr(depth_profile, "fallback_depth_m", 2.0))
|
||||
if num > 0:
|
||||
if np.any(line_mask):
|
||||
_, indices = ndimage.distance_transform_edt(~line_mask, return_indices=True)
|
||||
depth_nearest = depth_line[tuple(indices)]
|
||||
else:
|
||||
depth_nearest = np.zeros_like(depth_line, dtype=np.float32)
|
||||
for comp in range(1, num + 1):
|
||||
comp_mask = labels == comp
|
||||
if order_nearest is not None:
|
||||
comp_orders = order_nearest[comp_mask]
|
||||
comp_orders = comp_orders[comp_orders > 0.0]
|
||||
if comp_orders.size > 0:
|
||||
order_val = int(round(float(np.median(comp_orders))))
|
||||
depth_val = _order_depth(order_val, depth_profile)
|
||||
else:
|
||||
depth_val = None
|
||||
else:
|
||||
depth_val = None
|
||||
|
||||
line_components = set(np.unique(labels[line_mask]))
|
||||
line_components.discard(0)
|
||||
if line_components:
|
||||
line_components_arr = np.array(sorted(line_components))
|
||||
use_line = np.isin(labels, line_components_arr)
|
||||
depth_map[use_line] = depth_nearest[use_line] * mask[use_line]
|
||||
else:
|
||||
use_line = np.zeros_like(mask_bin, dtype=bool)
|
||||
if depth_val is None or depth_val <= 0.0:
|
||||
depth_val = fallback_depth
|
||||
|
||||
fallback_depth = float(getattr(depth_profile, "fallback_depth_m", 2.0))
|
||||
if fallback_depth > 0.0:
|
||||
fallback_mask = mask_bin & ~use_line
|
||||
depth_map[fallback_mask] = fallback_depth * mask[fallback_mask]
|
||||
else:
|
||||
use_line = np.zeros_like(mask_bin, dtype=bool)
|
||||
else:
|
||||
use_line = np.zeros_like(mask_bin, dtype=bool)
|
||||
if depth_val and depth_val > 0.0:
|
||||
depth_map[comp_mask] = depth_val * mask[comp_mask]
|
||||
|
||||
height_mosaic = height_mosaic - depth_map
|
||||
|
||||
flow_slope = float(getattr(depth_profile, "flow_slope_m_per_km", 0.0))
|
||||
dist_field = str(getattr(depth_profile, "flow_distance_field", "DIST_DN_KM") or "DIST_DN_KM")
|
||||
if flow_slope > 0.0 and np.any(line_mask) and indices is not None:
|
||||
if flow_slope > 0.0 and np.any(line_mask) and indices is not None and labels is not None and num > 0:
|
||||
dist_line = _build_flow_distance(
|
||||
layer_sources,
|
||||
mosaic_bounds,
|
||||
@@ -786,18 +782,24 @@ def erode_rivers(cfg: Config) -> int:
|
||||
tile_srs,
|
||||
)
|
||||
dist_nearest = dist_line[tuple(indices)]
|
||||
line_labels, line_num = ndimage.label(line_mask)
|
||||
if line_num > 0:
|
||||
flow_map = np.zeros_like(depth_map, dtype=np.float32)
|
||||
for comp in range(1, line_num + 1):
|
||||
comp_mask = line_labels == comp
|
||||
if not np.any(comp_mask):
|
||||
continue
|
||||
comp_max = float(np.max(dist_nearest[comp_mask]))
|
||||
flow_map[comp_mask] = (comp_max - dist_nearest[comp_mask]) * flow_slope
|
||||
depth_map[use_line] += flow_map[use_line] * mask[use_line]
|
||||
for comp in range(1, num + 1):
|
||||
comp_mask = labels == comp
|
||||
if not np.any(comp_mask):
|
||||
continue
|
||||
dist_vals = dist_nearest[comp_mask]
|
||||
heights = height_mosaic[comp_mask]
|
||||
if dist_vals.size == 0 or heights.size == 0:
|
||||
continue
|
||||
order = np.argsort(-dist_vals)
|
||||
dist_sorted = dist_vals[order]
|
||||
heights_sorted = heights[order]
|
||||
adjusted = heights_sorted + flow_slope * dist_sorted
|
||||
enforced = np.minimum.accumulate(adjusted)
|
||||
new_heights_sorted = enforced - flow_slope * dist_sorted
|
||||
new_heights = heights.copy()
|
||||
new_heights[order] = new_heights_sorted
|
||||
height_mosaic[comp_mask] = np.minimum(heights, new_heights)
|
||||
|
||||
height_mosaic = height_mosaic - depth_map
|
||||
|
||||
x0_center = (0 - min_dx) * step
|
||||
y0_center = (max_dy - 0) * step
|
||||
|
||||
Reference in New Issue
Block a user