building.py

"""Building isolation and cropping.

Selects the primary building from a multi-building mask using
geocoded coordinates, then crops imagery to the building footprint.
"""

import numpy as np
import cv2
from scipy.ndimage import label
from PIL import Image


def isolate_primary_building(
    building_mask: np.ndarray,
    center_lat: float,
    center_lng: float,
    bounds: tuple,
) -> np.ndarray | None:
    """Select the building closest to the geocoded address.

    Uses connected component labeling to identify separate buildings,
    then picks the one containing or nearest to the center coordinates.

    Args:
        building_mask: Binary mask (H, W) from Google Solar API.
        center_lat: Latitude from geocoding.
        center_lng: Longitude from geocoding.
        bounds: (west, south, east, north) in WGS84.

    Returns:
        Binary mask of just the primary building, or None.
    """
    if building_mask is None:
        return None

    labeled_mask, num_buildings = label(building_mask > 0)

    if num_buildings == 0:
        return None
    if num_buildings == 1:
        return (labeled_mask == 1).astype(np.uint8)

    # Convert lat/lng to pixel coordinates
    h, w = building_mask.shape
    west, south, east, north = bounds
    center_x = int((center_lng - west) / (east - west) * w)
    center_y = int((north - center_lat) / (north - south) * h)
    center_x = max(0, min(w - 1, center_x))
    center_y = max(0, min(h - 1, center_y))

    # Check if center is inside a building
    center_building_id = labeled_mask[center_y, center_x]
    if center_building_id > 0:
        return (labeled_mask == center_building_id).astype(np.uint8)

    # Find nearest building
    min_dist = float("inf")
    closest_id = 1
    for building_id in range(1, num_buildings + 1):
        pixels = np.argwhere(labeled_mask == building_id)
        if len(pixels) == 0:
            continue
        distances = np.sqrt(
            (pixels[:, 0] - center_y) ** 2 + (pixels[:, 1] - center_x) ** 2
        )
        d = distances.min()
        if d < min_dist:
            min_dist = d
            closest_id = building_id

    return (labeled_mask == closest_id).astype(np.uint8)


def crop_to_building(
    image: np.ndarray | Image.Image,
    building_mask: np.ndarray,
    padding_pct: float = 0.1,
) -> tuple[np.ndarray, dict, np.ndarray]:
    """Crop image and mask to building bounding box with padding.

    Args:
        image: RGB image (np.ndarray or PIL Image).
        building_mask: Binary mask (H, W).
        padding_pct: Padding as fraction of building size.

    Returns:
        (cropped_image_array, crop_info, cropped_mask)
        crop_info has keys: rmin, rmax, cmin, cmax, original_shape.
    """
    if isinstance(image, Image.Image):
        image = np.array(image)

    if building_mask is None:
        return image, None, None

    # Resize mask to match image if needed
    if building_mask.shape != image.shape[:2]:
        building_mask = cv2.resize(
            building_mask.astype(np.uint8),
            (image.shape[1], image.shape[0]),
            interpolation=cv2.INTER_NEAREST,
        )

    rows = np.any(building_mask > 0, axis=1)
    cols = np.any(building_mask > 0, axis=0)

    if not rows.any() or not cols.any():
        return image, None, None

    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]

    h, w = image.shape[:2]
    pad_h = int((rmax - rmin) * padding_pct)
    pad_w = int((cmax - cmin) * padding_pct)

    rmin = max(0, rmin - pad_h)
    rmax = min(h, rmax + pad_h)
    cmin = max(0, cmin - pad_w)
    cmax = min(w, cmax + pad_w)

    crop_info = {
        "rmin": rmin,
        "rmax": rmax,
        "cmin": cmin,
        "cmax": cmax,
        "original_shape": (h, w),
    }

    return image[rmin:rmax, cmin:cmax], crop_info, building_mask[rmin:rmax, cmin:cmax]


def recalculate_bounds(bounds: tuple, crop_info: dict) -> tuple:
    """Recalculate WGS84 bounds after cropping.

    Args:
        bounds: Original (west, south, east, north).
        crop_info: From crop_to_building.

    Returns:
        New (west, south, east, north) for the cropped region.
    """
    west, south, east, north = bounds
    orig_h, orig_w = crop_info["original_shape"]
    rmin, rmax = crop_info["rmin"], crop_info["rmax"]
    cmin, cmax = crop_info["cmin"], crop_info["cmax"]

    new_west = west + (east - west) * cmin / orig_w
    new_east = west + (east - west) * cmax / orig_w
    new_north = north - (north - south) * rmin / orig_h
    new_south = north - (north - south) * rmax / orig_h

    return (new_west, new_south, new_east, new_north)