Stream Threshold Tuning: Automating Network Extraction in Python

conda create -n hydro python=3.11
conda activate hydro
conda install -c conda-forge rasterio=1.3 numpy scipy geopandas shapely scikit-image whitebox richdem networkx

SCA = cell_count × cell_area_m²

A · S^n ≥ C_t

import logging
import rasterio
import numpy as np

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
log = logging.getLogger(__name__)

def load_accumulation(acc_path: str) -> tuple[np.ndarray, object, object, float]:
    """Load and validate a flow accumulation raster."""
    with rasterio.open(acc_path) as src:
        if not src.crs.is_projected:
            raise ValueError(
                f"Accumulation raster must use a projected CRS; got {src.crs}. "
                "Reproject to UTM or a local equal-area projection before thresholding."
            )
        cell_size_m = abs(src.res[0])
        acc = src.read(1).astype(np.float32)
        nodata = src.nodata
        if nodata is not None:
            acc[acc == nodata] = 0.0
        log.info(
            "Loaded accumulation: shape=%s, cell_size=%.2f m, CRS=%s",
            acc.shape, cell_size_m, src.crs,
        )
        return acc, src.transform, src.crs, cell_size_m

def accumulation_to_sca(acc: np.ndarray, cell_size_m: float) -> np.ndarray:
    """Convert cell-count accumulation to specific catchment area in m²."""
    return acc * (cell_size_m ** 2)

def apply_threshold(sca: np.ndarray, threshold_m2: float) -> np.ndarray:
    """Return binary channel mask for cells exceeding threshold SCA."""
    return (sca >= threshold_m2).astype(np.uint8)

import geopandas as gpd
from scipy import ndimage
from rasterio.features import rasterize

def score_threshold(
    sca: np.ndarray,
    threshold_m2: float,
    transform,
    crs,
    ref_buffered: np.ndarray,
    opening_iterations: int = 2,
) -> float:
    """
    Score a candidate threshold using the F1 metric.

    precision = TP / (TP + FP)
    recall    = TP / (TP + FN)
    F1        = 2 * precision * recall / (precision + recall)
    """
    structure = ndimage.generate_binary_structure(2, 1)
    mask = ndimage.binary_opening(
        sca >= threshold_m2, structure=structure, iterations=opening_iterations
    ).astype(np.uint8)

    tp = int(np.sum((mask == 1) & (ref_buffered == 1)))
    fp = int(np.sum((mask == 1) & (ref_buffered == 0)))
    fn = int(np.sum((mask == 0) & (ref_buffered == 1)))

    precision = tp / (tp + fp + 1e-9)
    recall    = tp / (tp + fn + 1e-9)
    f1 = 2 * precision * recall / (precision + recall + 1e-9)
    log.debug(
        "threshold=%.0f m²  P=%.3f  R=%.3f  F1=%.3f",
        threshold_m2, precision, recall, f1,
    )
    return f1


def sweep_thresholds(
    sca: np.ndarray,
    transform,
    crs,
    ref_gdf: gpd.GeoDataFrame,
    candidates: list[float] | None = None,
    buffer_m: float = 30.0,
    opening_iterations: int = 2,
) -> tuple[float, float]:
    """Return (best_threshold_m2, best_f1) from a sweep over candidate SCA values."""
    if candidates is None:
        # Logarithmic spacing from 5,000 m² to 500,000 m²
        candidates = list(np.geomspace(5_000, 500_000, num=20))

    ref_buffered = rasterize(
        [(geom, 1) for geom in ref_gdf.buffer(buffer_m).geometry],
        out_shape=sca.shape,
        transform=transform,
        fill=0,
        dtype=np.uint8,
    )

    best_t, best_f1 = candidates[0], 0.0
    for t in candidates:
        f1 = score_threshold(sca, t, transform, crs, ref_buffered, opening_iterations)
        if f1 > best_f1:
            best_f1, best_t = f1, t

    log.info("Best threshold: %.0f m²  (F1=%.3f)", best_t, best_f1)
    return best_t, best_f1

from shapely.geometry import LineString
from skimage.morphology import skeletonize as skimage_skeletonize

def vectorize_channels(
    sca: np.ndarray,
    threshold_m2: float,
    cell_size_m: float,
    transform,
    crs,
    min_segment_length_m: float = 30.0,
    opening_iterations: int = 2,
) -> gpd.GeoDataFrame:
    """
    Apply threshold, skeletonize to single-pixel-width centrelines,
    label connected components, and return a GeoDataFrame of LineStrings.
    Segments shorter than min_segment_length_m are discarded.
    """
    structure = ndimage.generate_binary_structure(2, 1)
    binary = ndimage.binary_opening(
        sca >= threshold_m2, structure=structure, iterations=opening_iterations
    )
    skeleton = skimage_skeletonize(binary).astype(np.uint8)

    labeled, n_features = ndimage.label(skeleton, structure=structure)
    log.info("Connected components after skeletonization: %d", n_features)

    lines = []
    for i in range(1, n_features + 1):
        coords_rc = np.argwhere(labeled == i)
        if len(coords_rc) * cell_size_m < min_segment_length_m:
            continue
        # Sort rows top-to-bottom as a crude downstream approximation for display
        coords_rc = coords_rc[np.argsort(coords_rc[:, 0])]
        xy = [rasterio.transform.xy(transform, r, c) for r, c in coords_rc]
        if len(xy) >= 2:
            lines.append(LineString(xy))

    gdf = gpd.GeoDataFrame(geometry=lines, crs=crs)
    log.info(
        "Extracted %d stream segments (min_length=%.0f m)",
        len(gdf), min_segment_length_m,
    )
    return gdf

def export_streams(gdf: gpd.GeoDataFrame, output_gpkg: str, layer: str = "streams") -> None:
    """Write stream network to a GeoPackage layer."""
    gdf.to_file(output_gpkg, driver="GPKG", layer=layer)
    log.info("Exported %d features to %s (layer=%s)", len(gdf), output_gpkg, layer)

import logging
import rasterio
import numpy as np
import geopandas as gpd
from scipy import ndimage
from skimage.morphology import skeletonize as skimage_skeletonize
from shapely.geometry import LineString
from rasterio.features import rasterize

log = logging.getLogger(__name__)


def tune_and_extract_streams(
    acc_path: str,
    ref_network_path: str,
    output_gpkg: str,
    candidates: list[float] | None = None,
    buffer_m: float = 30.0,
    min_segment_length_m: float = 30.0,
    opening_iterations: int = 2,
) -> gpd.GeoDataFrame:
    """
    End-to-end stream threshold tuning and extraction pipeline.

    Parameters
    ----------
    acc_path : path to flow accumulation raster (projected CRS, float32)
    ref_network_path : path to reference stream vector (same CRS)
    output_gpkg : path for exported GeoPackage
    candidates : list of SCA thresholds in m² to evaluate; defaults to 20
                 log-spaced values between 5,000 and 500,000 m²
    buffer_m : reference network buffer radius used in F1 scoring
    min_segment_length_m : discard extracted segments shorter than this
    opening_iterations : binary opening passes for noise suppression

    Returns
    -------
    GeoDataFrame of extracted stream LineStrings
    """
    # --- Load and validate ---
    with rasterio.open(acc_path) as src:
        if not src.crs.is_projected:
            raise ValueError(
                f"Accumulation raster CRS must be projected; got {src.crs}."
            )
        cell_size_m = abs(src.res[0])
        raw_acc = src.read(1).astype(np.float32)
        nodata = src.nodata
        transform = src.transform
        crs = src.crs
    if nodata is not None:
        raw_acc[raw_acc == nodata] = 0.0
    log.info(
        "Accumulation loaded: shape=%s  cell=%.1f m  CRS=%s",
        raw_acc.shape, cell_size_m, crs,
    )

    sca = raw_acc * (cell_size_m ** 2)
    log.info("SCA range: %.0f – %.0f m²", float(sca.min()), float(sca.max()))

    # --- Load reference network ---
    ref_gdf = gpd.read_file(ref_network_path).to_crs(crs)
    log.info("Reference network: %d features", len(ref_gdf))

    # --- Pre-rasterize reference buffer (reused across all sweep iterations) ---
    if candidates is None:
        candidates = list(np.geomspace(5_000, 500_000, num=20))

    ref_buffered = rasterize(
        [(geom, 1) for geom in ref_gdf.buffer(buffer_m).geometry],
        out_shape=sca.shape,
        transform=transform,
        fill=0,
        dtype=np.uint8,
    )

    structure = ndimage.generate_binary_structure(2, 1)
    best_t, best_f1 = candidates[0], 0.0
    for t in candidates:
        mask = ndimage.binary_opening(
            sca >= t, structure=structure, iterations=opening_iterations
        ).astype(np.uint8)
        tp = int(np.sum((mask == 1) & (ref_buffered == 1)))
        fp = int(np.sum((mask == 1) & (ref_buffered == 0)))
        fn = int(np.sum((mask == 0) & (ref_buffered == 1)))
        prec = tp / (tp + fp + 1e-9)
        rec  = tp / (tp + fn + 1e-9)
        f1   = 2 * prec * rec / (prec + rec + 1e-9)
        log.debug("  t=%.0f m²  P=%.3f  R=%.3f  F1=%.3f", t, prec, rec, f1)
        if f1 > best_f1:
            best_f1, best_t = f1, t
    log.info("Optimal threshold: %.0f m²  (F1=%.3f)", best_t, best_f1)

    # --- Apply optimal threshold and skeletonize ---
    binary = ndimage.binary_opening(
        sca >= best_t, structure=structure, iterations=opening_iterations
    )
    skeleton = skimage_skeletonize(binary).astype(np.uint8)
    labeled, n_comps = ndimage.label(skeleton, structure=structure)
    log.info("Connected components: %d", n_comps)

    lines = []
    for i in range(1, n_comps + 1):
        rc = np.argwhere(labeled == i)
        if len(rc) * cell_size_m < min_segment_length_m:
            continue
        rc = rc[np.argsort(rc[:, 0])]
        xy = [rasterio.transform.xy(transform, r, c) for r, c in rc]
        if len(xy) >= 2:
            lines.append(LineString(xy))

    gdf = gpd.GeoDataFrame(geometry=lines, crs=crs)
    log.info("Extracted %d segments (min_length=%.0f m)", len(gdf), min_segment_length_m)

    # --- Export ---
    gdf.to_file(output_gpkg, driver="GPKG", layer="streams")
    log.info("Saved to %s", output_gpkg)
    return gdf

import geopandas as gpd
from shapely.ops import unary_union

def compute_jaccard(
    extracted: gpd.GeoDataFrame,
    reference: gpd.GeoDataFrame,
    buffer_m: float = 30.0,
) -> float:
    """Compute Jaccard index between buffered extracted and reference networks."""
    ext_buf = unary_union(extracted.buffer(buffer_m))
    ref_buf = unary_union(reference.buffer(buffer_m))
    intersection = ext_buf.intersection(ref_buf).area
    union = ext_buf.union(ref_buf).area
    jaccard = intersection / union if union > 0 else 0.0
    log.info("Jaccard index: %.4f (buffer=%.0f m)", jaccard, buffer_m)
    return jaccard

import networkx as nx
from shapely.geometry import Point

def check_topology(gdf: gpd.GeoDataFrame) -> dict:
    """
    Build an undirected graph from segment endpoints and report
    connected-component count and any dangling single-node subgraphs.
    """
    G = nx.Graph()
    for geom in gdf.geometry:
        coords = list(geom.coords)
        start = (round(coords[0][0], 2), round(coords[0][1], 2))
        end   = (round(coords[-1][0], 2), round(coords[-1][1], 2))
        G.add_edge(start, end)

    components = list(nx.connected_components(G))
    dangling   = [c for c in components if len(c) == 1]
    log.info(
        "Graph components: %d total, %d dangling nodes",
        len(components), len(dangling),
    )
    return {"components": len(components), "dangling": len(dangling)}

def length_ratio(extracted: gpd.GeoDataFrame, reference: gpd.GeoDataFrame) -> float:
    """Ratio of total extracted length to total reference length."""
    ext_len = extracted.geometry.length.sum()
    ref_len = reference.geometry.length.sum()
    ratio = ext_len / ref_len if ref_len > 0 else float("inf")
    log.info("Length ratio: %.3f (extracted / reference)", ratio)
    return ratio