geo_export.py

"""GeoJSON polygon extraction and export.

Converts pixel-space label maps to georeferenced GeoJSON polygons
with area, pitch, and azimuth properties.
"""

import math
import numpy as np
import cv2


def pixel_to_geo(
    x: float, y: float, img_w: int, img_h: int, bounds: tuple
) -> list[float]:
    """Convert pixel coordinates to [longitude, latitude] (GeoJSON order).

    Y-axis is flipped: pixel y=0 is north, y=max is south.
    """
    west, south, east, north = bounds
    lng = west + (east - west) * (x / img_w)
    lat = north - (north - south) * (y / img_h)
    return [lng, lat]


def geo_to_pixel(
    lng: float, lat: float, img_w: int, img_h: int, bounds: tuple
) -> tuple[int, int]:
    """Convert geographic coordinates to pixel coordinates."""
    west, south, east, north = bounds
    px = int((lng - west) / (east - west) * img_w)
    py = int((north - lat) / (north - south) * img_h)
    return px, py


def _compute_area(contour_area_px: float, img_w: int, img_h: int, bounds: tuple) -> tuple[float, float]:
    """Compute area in sqm and sqft from pixel area."""
    west, south, east, north = bounds
    center_lat = (north + south) / 2
    meters_per_lng = 111320 * math.cos(math.radians(center_lat))
    meters_per_lat = 111320
    pixel_width_m = (east - west) * meters_per_lng / img_w
    pixel_height_m = (north - south) * meters_per_lat / img_h
    area_sqm = contour_area_px * pixel_width_m * pixel_height_m
    area_sqft = area_sqm * 10.764
    return area_sqm, area_sqft


def _is_boundary_rect(contour, img_w: int, img_h: int, tolerance: int = 5) -> bool:
    """Check if contour matches the full image bounds (false positive)."""
    epsilon = 0.02 * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, epsilon, True)
    if len(approx) != 4:
        return False

    x, y, w, h = cv2.boundingRect(contour)
    return (
        x < tolerance
        and y < tolerance
        and (x + w) > (img_w - tolerance)
        and (y + h) > (img_h - tolerance)
    )


def labels_to_geojson(
    label_map: np.ndarray,
    bounds: tuple,
    plane_info: list[dict],
    min_area_sqft: float = 50.0,
) -> dict:
    """Convert a label map to a GeoJSON FeatureCollection.

    Args:
        label_map: Integer label map (H, W). 0 = background.
        bounds: (west, south, east, north) in WGS84.
        plane_info: List of dicts with plane metadata (keyed by segment_id).
        min_area_sqft: Minimum polygon area to include.

    Returns:
        GeoJSON FeatureCollection dict.
    """
    h, w = label_map.shape
    features = []

    # Build lookup from segment_id to plane info
    info_lookup = {p["segment_id"]: p for p in plane_info}

    unique_labels = sorted(set(np.unique(label_map)) - {0})

    for label_id in unique_labels:
        mask = (label_map == label_id).astype(np.uint8)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for contour in contours:
            px_area = cv2.contourArea(contour)
            if px_area < 100:
                continue
            if _is_boundary_rect(contour, w, h):
                continue

            area_sqm, area_sqft = _compute_area(px_area, w, h, bounds)
            if area_sqft < min_area_sqft:
                continue

            # Simplify polygon (2.5% of perimeter — clean roof plane shapes)
            perimeter = cv2.arcLength(contour, True)
            epsilon = 0.025 * perimeter
            simplified = cv2.approxPolyDP(contour, epsilon, True)

            coords = [pixel_to_geo(pt[0][0], pt[0][1], w, h, bounds) for pt in simplified]

            # Close the ring
            if coords and coords[0] != coords[-1]:
                coords.append(coords[0])

            if len(coords) < 4:
                continue

            # Merge plane properties
            props = {
                "roof_plane_id": int(label_id),
                "area_sqm": round(area_sqm, 2),
                "area_sqft": round(area_sqft, 2),
                "num_vertices": len(coords) - 1,
            }
            info = info_lookup.get(label_id, {})
            if "mean_slope" in info:
                props["pitch_deg"] = round(info["mean_slope"], 1)
            if "mean_aspect" in info:
                props["azimuth_deg"] = round(info["mean_aspect"], 1)
            if "mean_height" in info:
                props["height_m"] = round(info["mean_height"], 2)

            features.append(
                {
                    "type": "Feature",
                    "properties": props,
                    "geometry": {"type": "Polygon", "coordinates": [coords]},
                }
            )

    return {
        "type": "FeatureCollection",
        "features": features,
    }