src/extract_centerline.py

"""
Digitize a hand-traced river line (a colored polyline drawn over a road
map, exported as a PNG/screenshot) into a georeferenced centerline CSV
of [longitude, latitude] waypoints, ordered along the river's course.

Pipeline:
  1. Color-threshold the image to isolate just the traced line.
  2. Keep only the largest connected component (drops stray marker dots
     etc. that happen to be a similar color).
  3. Skeletonize to a 1px-wide path.
  4. Find the true end-to-end path via "double BFS" (find the farthest
     node from an arbitrary start, then the farthest node from THAT one
     — the shortest path between those two is the main line; this
     naturally ignores short skeletonization artifacts/spurs without
     needing to hand-tune a spur-pruning threshold).
  5. Georeference: fit a least-squares affine transform (pixel -> lon/lat)
     from a handful of reference points (town positions you read off the
     image, paired with their known real-world coordinates).
  6. Downsample to a manageable number of evenly-spaced waypoints and
     save as CSV.

Usage (edit REFERENCE_POINTS for your image, or pass --refs-json):
    python extract_river_centerline.py \\
        --image map.png --target-rgb 30 144 255 --tolerance 40 \\
        --refs-json refs.json --output centerline.csv --n-points 120

Where refs.json looks like:
    [
      {"pixel_x": 200, "pixel_y": 45, "lon": 1.1667, "lat": 49.2167},
      {"pixel_x": 200, "pixel_y": 183, "lon": 1.1509, "lat": 49.0269},
      ...
    ]

ACCURACY NOTE: this only produces a shape that's as good as (a) the
traced line's own accuracy and (b) how precisely you read reference
point pixel coordinates off the image. Expect a few km of error, not
survey-grade. Use >= 4-5 well-spread reference points for a decent
least-squares fit; more points and points that span the image well
help.

Requires: numpy, pandas, Pillow, scikit-image, networkx
"""
import argparse
import json
from pathlib import Path
from typing import List, Dict

import numpy as np
import pandas as pd
from PIL import Image
from skimage.morphology import skeletonize
from skimage.measure import label
import networkx as nx


def isolate_line_mask(image_path: Path, target_rgb: List[int], tolerance: float) -> np.ndarray:
    """Threshold the image to a boolean mask of pixels close to `target_rgb`."""
    img = Image.open(image_path).convert("RGB")
    arr = np.array(img).astype(int)
    dist = np.sqrt(((arr - np.array(target_rgb)) ** 2).sum(axis=2))
    return dist < tolerance


def keep_largest_component(mask: np.ndarray) -> np.ndarray:
    """Drop everything except the largest connected blob (e.g. excludes a
    stray start/end marker dot rendered in a similar color)."""
    lbl = label(mask)
    if lbl.max() == 0:
        raise ValueError("No pixels matched the target color — check target_rgb/tolerance.")
    sizes = [((lbl == i).sum(), i) for i in range(1, lbl.max() + 1)]
    largest_id = max(sizes)[1]
    return lbl == largest_id


def order_skeleton_pixels(skel: np.ndarray) -> np.ndarray:
    """
    Skeletonize-then-order: build an 8-connectivity graph over skeleton
    pixels, then find the path between the two most mutually-distant
    nodes (double BFS). Returns an (N, 2) array of (y, x) pixel
    coordinates in order along the line.
    """
    ys, xs = np.where(skel)
    coords = set(zip(ys.tolist(), xs.tolist()))

    G = nx.Graph()
    G.add_nodes_from(coords)
    for (y, x) in coords:
        for dy in (-1, 0, 1):
            for dx in (-1, 0, 1):
                if dy == 0 and dx == 0:
                    continue
                q = (y + dy, x + dx)
                if q in coords:
                    G.add_edge((y, x), q)

    if nx.number_connected_components(G) > 1:
        # Keep only the largest component of the skeleton graph too.
        largest_cc = max(nx.connected_components(G), key=len)
        G = G.subgraph(largest_cc).copy()

    start = next(iter(G.nodes))
    lengths = nx.single_source_shortest_path_length(G, start)
    a = max(lengths, key=lengths.get)
    lengths2 = nx.single_source_shortest_path_length(G, a)
    b = max(lengths2, key=lengths2.get)

    path = nx.shortest_path(G, a, b)
    return np.array(path)  # (N, 2) as (y, x)


def fit_pixel_to_geo_transform(refs: List[Dict[str, float]]):
    """
    Least-squares affine fit: lon = a*px + b*py + c, lat = d*px + e*py + f.
    Returns (lon_coef, lat_coef), each a length-3 array [a, b, c].
    Also prints residuals at the reference points as a sanity check.
    """
    P = np.array([[r["pixel_x"], r["pixel_y"], 1] for r in refs])
    lon = np.array([r["lon"] for r in refs])
    lat = np.array([r["lat"] for r in refs])

    lon_coef, *_ = np.linalg.lstsq(P, lon, rcond=None)
    lat_coef, *_ = np.linalg.lstsq(P, lat, rcond=None)

    pred_lon, pred_lat = P @ lon_coef, P @ lat_coef
    mean_lat = np.mean(lat)
    err_lon_km = np.abs(pred_lon - lon) * 111 * np.cos(np.radians(mean_lat))
    err_lat_km = np.abs(pred_lat - lat) * 111
    print(f"Georeferencing fit residuals at the {len(refs)} reference points: "
          f"max {max(err_lon_km.max(), err_lat_km.max()):.2f} km "
          f"(lon max {err_lon_km.max():.2f} km, lat max {err_lat_km.max():.2f} km)")

    return lon_coef, lat_coef


def pixels_to_geo(path_px: np.ndarray, lon_coef: np.ndarray, lat_coef: np.ndarray) -> pd.DataFrame:
    """Apply the fitted transform to an (N, 2) array of (y, x) pixel coords."""
    ys, xs = path_px[:, 0], path_px[:, 1]
    lon = lon_coef[0] * xs + lon_coef[1] * ys + lon_coef[2]
    lat = lat_coef[0] * xs + lat_coef[1] * ys + lat_coef[2]
    return pd.DataFrame({"longitude": lon, "latitude": lat})


def downsample_evenly(df: pd.DataFrame, n_points: int) -> pd.DataFrame:
    """Evenly-spaced-by-index downsampling (fine since the ordered path's
    pixels are already roughly evenly spaced along the traced curve)."""
    idx = np.linspace(0, len(df) - 1, n_points).astype(int)
    out = df.iloc[idx].reset_index(drop=True)
    out.insert(0, "seq", range(len(out)))
    return out


def extract_centerline(
    image_path: Path,
    target_rgb: List[int],
    tolerance: float,
    reference_points: List[Dict[str, float]],
    n_points: int = 120,
) -> pd.DataFrame:
    """Run the full pipeline and return the final waypoints DataFrame."""
    mask = isolate_line_mask(image_path, target_rgb, tolerance)
    mask = keep_largest_component(mask)
    print(f"Isolated line: {mask.sum()} pixels")

    skel = skeletonize(mask)
    print(f"Skeletonized: {skel.sum()} pixels")

    path_px = order_skeleton_pixels(skel)
    coverage = len(path_px) / max(skel.sum(), 1) * 100
    print(f"Ordered path: {len(path_px)} pixels ({coverage:.0f}% of the skeleton — "
          f"a low % here suggests a branchy/noisy skeleton worth checking visually, "
          f"e.g. by plotting the ordered path over the original image)")

    lon_coef, lat_coef = fit_pixel_to_geo_transform(reference_points)
    geo_df = pixels_to_geo(path_px, lon_coef, lat_coef)

    result = downsample_evenly(geo_df, n_points)
    print(f"Final centerline: {len(result)} waypoints")
    return result


# ----------------------------------------------------------------------
# Reference points actually used for La Eure and La Risle in this project.
# Read by eye off a pixel-gridded version of each map image, paired with
# each town's known real-world coordinates. Kept here for reproducibility
# and as a template for digitizing further rivers.
# ----------------------------------------------------------------------
EURE_REFERENCE_POINTS = [
    {"pixel_x": 200, "pixel_y": 45, "lon": 1.1667, "lat": 49.2167},   # Louviers
    {"pixel_x": 200, "pixel_y": 183, "lon": 1.1509, "lat": 49.0269}, # Evreux
    {"pixel_x": 290, "pixel_y": 325, "lon": 1.3625, "lat": 48.7377}, # Dreux
    {"pixel_x": 325, "pixel_y": 125, "lon": 1.4836, "lat": 49.0917}, # Vernon
    {"pixel_x": 295, "pixel_y": 495, "lon": 1.4894, "lat": 48.4439}, # Chartres
    {"pixel_x": 300, "pixel_y": 88, "lon": 1.3333, "lat": 49.1667},  # Gaillon
    {"pixel_x": 90, "pixel_y": 550, "lon": 0.8244, "lat": 48.3236},  # Nogent-le-Rotrou
]
EURE_TARGET_RGB = [30, 144, 255]  # "dodgerblue"-ish line color used in that map export

RISLE_REFERENCE_POINTS = [
    {"pixel_x": 55, "pixel_y": 42, "lon": 0.2333, "lat": 49.4197},   # Honfleur
    {"pixel_x": 175, "pixel_y": 78, "lon": 0.5136, "lat": 49.3547},  # Pont-Audemer
    {"pixel_x": 245, "pixel_y": 220, "lon": 0.7169, "lat": 49.1936}, # Brionne
    {"pixel_x": 178, "pixel_y": 250, "lon": 0.5981, "lat": 49.0906}, # Bernay
    {"pixel_x": 195, "pixel_y": 415, "lon": 0.6339, "lat": 48.7642}, # L'Aigle
    {"pixel_x": 55, "pixel_y": 200, "lon": 0.2306, "lat": 49.1447},  # Lisieux
]
RISLE_TARGET_RGB = [0, 0, 205]  # "mediumblue"-ish line color used in that map export


def main():
    parser = argparse.ArgumentParser(description="Digitize a traced river line into a centerline CSV")
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--target-rgb", type=int, nargs=3, required=True, metavar=("R", "G", "B"),
                         help="RGB color of the traced line to isolate")
    parser.add_argument("--tolerance", type=float, default=40.0,
                         help="Color-distance tolerance for isolating the line (default: 40)")
    parser.add_argument("--refs-json", type=Path, default=None,
                         help="JSON file of reference points [{pixel_x, pixel_y, lon, lat}, ...]. "
                              "If omitted, --preset must be used instead.")
    parser.add_argument("--preset", choices=["eure", "risle"], default=None,
                         help="Use the built-in reference points for La Eure or La Risle "
                              "(the ones used in this project) instead of --refs-json")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--n-points", type=int, default=120)
    args = parser.parse_args()

    if args.preset == "eure":
        refs = EURE_REFERENCE_POINTS
    elif args.preset == "risle":
        refs = RISLE_REFERENCE_POINTS
    elif args.refs_json:
        refs = json.loads(args.refs_json.read_text())
    else:
        parser.error("Provide either --refs-json or --preset eure|risle")

    result = extract_centerline(args.image, args.target_rgb, args.tolerance, refs, args.n_points)
    result.to_csv(args.output, index=False)
    print(f"Saved to {args.output}")


if __name__ == "__main__":
    main()