diff --git a/geodata_pipeline/river_erosion.py b/geodata_pipeline/river_erosion.py index da308ec..84348c1 100644 --- a/geodata_pipeline/river_erosion.py +++ b/geodata_pipeline/river_erosion.py @@ -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