ransac_planes.py

"""RANSAC 3D plane fitting for DSM-based roof segmentation.

Fits geometric planes to the Digital Surface Model point cloud,
producing clean plane boundaries that are completely shadow-invariant.
"""

import numpy as np
import cv2
from scipy.ndimage import distance_transform_edt

PIXEL_SIZE_M = 0.1  # Google Solar API resolution (10 cm/pixel)


def preprocess_dsm(dsm: np.ndarray, building_mask: np.ndarray | None = None) -> np.ndarray:
    """Smooth DSM while preserving ridge edges.

    1. Median filter (3x3) removes outlier spikes.
    2. Bilateral filter preserves edges while smoothing flat areas.
    """
    dsm_clean = cv2.medianBlur(dsm.astype(np.float32), 3)

    height_range = dsm_clean.max() - dsm_clean.min()
    if height_range < 1e-6:
        return dsm_clean

    # Bilateral filter operates on uint8, so normalize
    dsm_norm = ((dsm_clean - dsm_clean.min()) / height_range * 255).astype(np.uint8)
    dsm_filtered = cv2.bilateralFilter(dsm_norm, d=9, sigmaColor=50, sigmaSpace=9)

    # Denormalize back to meters
    dsm_smooth = (dsm_filtered.astype(np.float32) / 255.0) * height_range + dsm_clean.min()
    return dsm_smooth


def dsm_to_point_cloud(
    dsm: np.ndarray,
    building_mask: np.ndarray,
    pixel_size: float = PIXEL_SIZE_M,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Convert DSM pixels within building mask to a 3D point cloud.

    Args:
        dsm: Height map (H, W) in meters.
        building_mask: Binary mask (H, W).
        pixel_size: Meters per pixel.

    Returns:
        (points_3d: Nx3, valid_flat_indices: N, valid_mask_flat: M bool)
        points_3d columns are [x_meters, y_meters, height].
    """
    h, w = dsm.shape

    if building_mask.shape != (h, w):
        building_mask = cv2.resize(
            building_mask.astype(np.uint8), (w, h),
            interpolation=cv2.INTER_NEAREST,
        )

    valid_mask = building_mask > 0
    valid_flat = valid_mask.flatten()

    yy, xx = np.mgrid[0:h, 0:w]
    xx_m = (xx * pixel_size).flatten()[valid_flat]
    yy_m = (yy * pixel_size).flatten()[valid_flat]
    zz = dsm.flatten()[valid_flat]

    points_3d = np.column_stack([xx_m, yy_m, zz])
    valid_flat_indices = np.where(valid_flat)[0]

    return points_3d, valid_flat_indices, valid_flat


def fit_planes(
    points_3d: np.ndarray,
    distance_threshold: float = 0.15,
    min_points: int = 500,
    max_planes: int = 8,
    n_iterations: int = 1000,
) -> list[dict]:
    """Iterative RANSAC plane fitting with SVD refinement.

    For each iteration:
      1. Sample 3 random points, compute plane via cross product.
      2. Count inliers within distance_threshold.
      3. Keep best hypothesis.
      4. Refit with SVD least-squares on all inliers.
      5. Remove inliers, repeat.

    Returns:
        List of plane dicts with keys: equation, normal, inlier_indices,
        pitch_deg, azimuth_deg, mean_height, num_points.
    """
    remaining_mask = np.ones(len(points_3d), dtype=bool)
    planes = []

    for plane_idx in range(max_planes):
        remaining_indices = np.where(remaining_mask)[0]
        remaining_pts = points_3d[remaining_indices]

        if len(remaining_pts) < min_points:
            break

        # RANSAC search
        best_count = 0
        best_inliers = None

        for _ in range(n_iterations):
            if len(remaining_pts) < 3:
                break

            idx = np.random.choice(len(remaining_pts), 3, replace=False)
            p1, p2, p3 = remaining_pts[idx]

            normal = np.cross(p2 - p1, p3 - p1)
            mag = np.linalg.norm(normal)
            if mag < 1e-8:
                continue

            normal /= mag
            d = -np.dot(normal, p1)

            distances = np.abs(remaining_pts @ normal + d)
            inliers = distances < distance_threshold
            count = inliers.sum()

            if count > best_count:
                best_count = count
                best_inliers = inliers

        if best_count < min_points:
            break

        # SVD refinement
        inlier_pts = remaining_pts[best_inliers]
        centroid = inlier_pts.mean(axis=0)
        _, _, vh = np.linalg.svd(inlier_pts - centroid)
        refined_normal = vh[2]
        refined_d = -np.dot(refined_normal, centroid)

        # Re-select inliers with refined plane
        distances = np.abs(remaining_pts @ refined_normal + refined_d)
        final_inliers = distances < distance_threshold

        # Map back to global indices
        global_inlier_idx = remaining_indices[final_inliers]
        remaining_mask[global_inlier_idx] = False

        # Ensure normal points upward
        if refined_normal[2] < 0:
            refined_normal = -refined_normal
            refined_d = -refined_d

        # Plane metrics
        pitch_deg = float(np.degrees(np.arccos(np.clip(abs(refined_normal[2]), 0, 1))))
        azimuth_deg = float(np.degrees(np.arctan2(-refined_normal[0], -refined_normal[1])) % 360)

        # Reject wall-like fits — roofs top out around 60° pitch (12/12 = 45°).
        # Anything steeper is almost certainly RANSAC finding a facade edge in the DSM.
        if pitch_deg > 65.0:
            # Still consume the inliers so we don't re-fit the same wall,
            # but don't emit this plane.
            continue

        planes.append({
            "segment_id": len(planes) + 1,
            "equation": [float(refined_normal[0]), float(refined_normal[1]),
                         float(refined_normal[2]), float(refined_d)],
            "normal": refined_normal.tolist(),
            "plane_d": float(refined_d),
            "inlier_indices": global_inlier_idx,
            "pitch_deg": pitch_deg,
            "azimuth_deg": azimuth_deg,
            "mean_height": float(points_3d[global_inlier_idx, 2].mean()),
            "num_points": int(len(global_inlier_idx)),
        })

    return planes


def planes_to_label_map(
    planes: list[dict],
    valid_flat_indices: np.ndarray,
    img_shape: tuple[int, int],
    building_mask: np.ndarray,
) -> np.ndarray:
    """Convert RANSAC plane results to a pixel label map.

    Assigns each building pixel to its plane. Unlabeled pixels
    are filled via nearest-neighbor distance transform.

    Returns:
        Label map (H, W) int32. 0 = background.
    """
    h, w = img_shape
    segments = np.zeros((h, w), dtype=np.int32)

    for plane in planes:
        global_idx = plane["inlier_indices"]
        flat_idx = valid_flat_indices[global_idx]
        py = flat_idx // w
        px = flat_idx % w
        segments[py, px] = plane["segment_id"]

    # Resize building mask if needed
    if building_mask.shape != (h, w):
        building_mask = cv2.resize(
            building_mask.astype(np.uint8), (w, h),
            interpolation=cv2.INTER_NEAREST,
        )

    # Fill gaps within building via nearest labeled pixel
    valid = building_mask > 0
    unlabeled = (segments == 0) & valid
    if unlabeled.any() and (segments > 0).any():
        _, nearest_idx = distance_transform_edt(unlabeled, return_indices=True)
        segments[unlabeled] = segments[nearest_idx[0][unlabeled], nearest_idx[1][unlabeled]]

    return segments


def build_plane_info(planes: list[dict], pixel_size: float = PIXEL_SIZE_M) -> list[dict]:
    """Build clean plane info list (without large inlier arrays).

    Returns list of dicts suitable for GeoJSON properties and display.
    """
    return [
        {
            "segment_id": p["segment_id"],
            "mean_slope": p["pitch_deg"],
            "mean_aspect": p["azimuth_deg"],
            "mean_height": p["mean_height"],
            "area_pixels": p["num_points"],
            "area_sqm": float(p["num_points"] * pixel_size ** 2),
            "plane_normal": p["normal"],
            "plane_d": p["plane_d"],
        }
        for p in planes
    ]


def upscale_dsm_with_planes(
    dsm: np.ndarray,
    planes: list[dict],
    label_map: np.ndarray,
    factor: int = 4,
    pixel_size: float = PIXEL_SIZE_M,
) -> np.ndarray:
    """Re-rasterize DSM at higher resolution using fitted plane equations.

    For each pixel assigned to a plane, compute the exact height from
    z = -(Ax + By + D) / C. Unassigned pixels get bilinear interpolation.

    Returns:
        High-resolution DSM array (H*factor, W*factor).
    """
    h, w = dsm.shape
    h_hr, w_hr = h * factor, w * factor

    # Build lookup: segment_id -> plane equation
    eq_lookup = {}
    for p in planes:
        eq_lookup[p["segment_id"]] = p["equation"]

    # High-res pixel grid in meters
    yy_hr, xx_hr = np.mgrid[0:h_hr, 0:w_hr]
    xx_m = xx_hr * (pixel_size / factor)
    yy_m = yy_hr * (pixel_size / factor)

    # Upscale label map to high-res
    label_hr = cv2.resize(label_map.astype(np.float32), (w_hr, h_hr),
                          interpolation=cv2.INTER_NEAREST).astype(np.int32)

    dsm_hr = np.full((h_hr, w_hr), np.nan, dtype=np.float32)

    for seg_id, (a, b, c, d) in eq_lookup.items():
        if abs(c) < 1e-10:
            continue
        mask = label_hr == seg_id
        dsm_hr[mask] = -(a * xx_m[mask] + b * yy_m[mask] + d) / c

    # Fill remaining pixels with bilinear interpolation of original
    from scipy.interpolate import RegularGridInterpolator

    unfilled = np.isnan(dsm_hr)
    if unfilled.any():
        row_coords = np.arange(h)
        col_coords = np.arange(w)
        interp = RegularGridInterpolator(
            (row_coords, col_coords), dsm.astype(np.float32),
            method="linear", bounds_error=False, fill_value=np.nan,
        )
        lr_rows = yy_hr[unfilled] / factor
        lr_cols = xx_hr[unfilled] / factor
        dsm_hr[unfilled] = interp(np.column_stack([lr_rows, lr_cols]))

    return dsm_hr