backend/app/modules/floor_plane.py

"""Floor-plane + perspective estimation (spec section 5, stage 2).

Combines the floor mask (SAM or geometric) with a simple pinhole-style ground
model. Full 3D reconstruction is intentionally OUT — a horizon estimate plus a
row-based scale function is sufficient for MVP placement.
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

from ..config import settings
from . import segmentation


@dataclass
class Perspective:
    """Row-based ground-plane scale model.

    ``ppc(y)`` gives lateral pixels-per-cm for an object resting on the floor at
    image row ``y``: large near the camera (image bottom), shrinking to ~0 toward
    the horizon. This is the heuristic used when no metric depth is available.
    """

    width: int
    height: int
    horizon_y: float
    floor_width_cm_bottom: float
    foreshorten: float

    def ppc(self, y: float) -> float:
        denom = max(1.0, self.height - self.horizon_y)
        t = (y - self.horizon_y) / denom         # 0 at horizon, 1 at bottom
        t = max(0.02, min(1.0, t))               # clamp away from zero
        return (self.width / self.floor_width_cm_bottom) * t


@dataclass
class FloorPlane:
    perspective: Perspective
    polygon: list[list[float]]   # floor polygon [[x, y], ...] in pixels
    mask: np.ndarray             # HxW uint8 (255 = floor)


def estimate_floor_plane(
    image_path: Path, floor_mask: np.ndarray | None = None
) -> FloorPlane:
    if floor_mask is None:
        floor_mask = segmentation.floor_mask(image_path)
    h, w = floor_mask.shape[:2]

    # Horizon ~ the highest floor row (where floor meets the back wall).
    rows_with_floor = np.where(floor_mask.max(axis=1) > 0)[0]
    horizon_y = float(rows_with_floor.min()) if rows_with_floor.size else h * settings.HORIZON_FRAC

    persp = Perspective(
        width=w,
        height=h,
        horizon_y=horizon_y,
        floor_width_cm_bottom=settings.ROOM_FLOOR_WIDTH_CM,
        foreshorten=settings.DEPTH_FORESHORTEN,
    )
    return FloorPlane(perspective=persp, polygon=_mask_to_polygon(floor_mask), mask=floor_mask)


def _mask_to_polygon(mask: np.ndarray, max_points: int = 24) -> list[list[float]]:
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        h, w = mask.shape[:2]
        return [[0, h - 1], [w - 1, h - 1], [w - 1, 0], [0, 0]]
    c = max(contours, key=cv2.contourArea)
    eps = 0.01 * cv2.arcLength(c, True)
    approx = cv2.approxPolyDP(c, eps, True).reshape(-1, 2)
    if len(approx) > max_points:
        idx = np.linspace(0, len(approx) - 1, max_points).astype(int)
        approx = approx[idx]
    return [[float(x), float(y)] for x, y in approx]