Upload floorplan/geometry.py

floorplan/geometry.py

@@ -0,0 +1,331 @@
+"""
+Geometry engine — compute wall polygons, detect rooms, and handle centerline operations.
+
+All operations use Shapely for robust computational geometry.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import numpy as np
+from shapely import ops
+from shapely.geometry import (
+    LineString,
+    MultiLineString,
+    MultiPolygon,
+    Point,
+    Polygon,
+)
+from shapely.ops import polygonize, unary_union
+
+from .schema import FloorPlan, Room, RoomLabel, Wall, WallFaceRef, WallSide, Point2D
+
+
+# ──────────────────────────────────────────────
+# Centerline utilities
+# ──────────────────────────────────────────────
+
+def compute_centerline_length(coords: list[tuple[float, float]]) -> float:
+    """Total length of a polyline defined by coordinate pairs."""
+    return LineString(coords).length
+
+
+def interpolate_along_centerline(
+    coords: list[tuple[float, float]],
+    distance: float,
+) -> tuple[float, float]:
+    """Find the point at a given distance along a centerline polyline."""
+    line = LineString(coords)
+    pt = line.interpolate(distance)
+    return (pt.x, pt.y)
+
+
+def normal_at_distance(
+    coords: list[tuple[float, float]],
+    distance: float,
+) -> tuple[float, float]:
+    """Unit normal vector (pointing LEFT of travel direction) at a distance along centerline."""
+    line = LineString(coords)
+    eps = min(0.001, line.length * 0.001)
+
+    d0 = max(0, distance - eps)
+    d1 = min(line.length, distance + eps)
+    p0 = line.interpolate(d0)
+    p1 = line.interpolate(d1)
+
+    dx = p1.x - p0.x
+    dy = p1.y - p0.y
+    length = math.hypot(dx, dy)
+    if length < 1e-12:
+        return (0.0, 1.0)
+
+    # Left-pointing normal: rotate direction 90° CCW
+    nx = -dy / length
+    ny = dx / length
+    return (nx, ny)
+
+
+def centerline_to_linestring(wall: Wall) -> LineString:
+    """Convert a Wall's centerline to a Shapely LineString."""
+    return LineString(wall.centerline_coords)
+
+
+# ──────────────────────────────────────────────
+# Wall polygon computation
+# ──────────────────────────────────────────────
+
+def wall_to_polygon(wall: Wall) -> Polygon:
+    """Compute the thick wall polygon by buffering the centerline.
+    
+    Uses flat end caps (cap_style=2) so wall ends are square-cut,
+    enabling clean joins with adjacent walls.
+    """
+    line = centerline_to_linestring(wall)
+    half_t = wall.thickness / 2.0
+    # cap_style: 1=round, 2=flat, 3=square
+    # join_style: 1=round, 2=mitre, 3=bevel
+    poly = line.buffer(half_t, cap_style=2, join_style=2)
+    if not poly.is_valid:
+        poly = poly.buffer(0)  # fix self-intersections
+    return poly
+
+
+def wall_to_left_right_lines(wall: Wall) -> tuple[LineString, LineString]:
+    """Compute the left and right offset lines of a wall.
+    
+    Left/right determined by walking along centerline from first to last point.
+    Left = your left side, Right = your right side.
+    
+    Returns (left_line, right_line).
+    """
+    line = centerline_to_linestring(wall)
+    half_t = wall.thickness / 2.0
+
+    left_line = line.parallel_offset(half_t, side="left")
+    right_line = line.parallel_offset(half_t, side="right")
+
+    # Shapely's parallel_offset reverses direction for right side — fix it
+    if isinstance(right_line, LineString) and len(right_line.coords) >= 2:
+        # Check if right_line goes in same direction as centerline
+        cl_start = np.array(line.coords[0])
+        r_start = np.array(right_line.coords[0])
+        r_end = np.array(right_line.coords[-1])
+        if np.linalg.norm(r_end - cl_start) < np.linalg.norm(r_start - cl_start):
+            right_line = LineString(list(right_line.coords)[::-1])
+
+    return (left_line, right_line)
+
+
+def wall_union(walls: list[Wall]) -> Polygon | MultiPolygon:
+    """Union of all wall polygons."""
+    polys = [wall_to_polygon(w) for w in walls]
+    return unary_union(polys)
+
+
+# ──────────────────────────────────────────────
+# Opening geometry
+# ──────────────────────────────────────────────
+
+def opening_to_gap_polygon(wall: Wall, opening_idx: int) -> Polygon:
+    """Compute the polygon of a door/window gap cut through a wall.
+    
+    The gap is a rectangle perpendicular to the wall at the opening location,
+    spanning the full wall thickness.
+    """
+    opening = wall.openings[opening_idx]
+    coords = wall.centerline_coords
+    half_t = wall.thickness / 2.0
+
+    # Get start and end points along centerline
+    start_pt = interpolate_along_centerline(coords, opening.start)
+    end_pt = interpolate_along_centerline(coords, opening.start + opening.length)
+
+    # Get normals at start and end
+    n_start = normal_at_distance(coords, opening.start)
+    n_end = normal_at_distance(coords, opening.start + opening.length)
+
+    # Build the gap rectangle: offset start/end points by ±half_thickness along normal
+    margin = half_t * 1.1  # slight margin for clean boolean ops
+    p1 = (start_pt[0] + n_start[0] * margin, start_pt[1] + n_start[1] * margin)
+    p2 = (end_pt[0] + n_end[0] * margin, end_pt[1] + n_end[1] * margin)
+    p3 = (end_pt[0] - n_end[0] * margin, end_pt[1] - n_end[1] * margin)
+    p4 = (start_pt[0] - n_start[0] * margin, start_pt[1] - n_start[1] * margin)
+
+    return Polygon([p1, p2, p3, p4])
+
+
+def wall_polygon_with_openings(wall: Wall) -> Polygon | MultiPolygon:
+    """Wall polygon with opening gaps cut out."""
+    poly = wall_to_polygon(wall)
+    for i in range(len(wall.openings)):
+        gap = opening_to_gap_polygon(wall, i)
+        poly = poly.difference(gap)
+        if not poly.is_valid:
+            poly = poly.buffer(0)
+    return poly
+
+
+# ──────────────────────────────────────────────
+# Room detection
+# ──────────────────────────────────────────────
+
+def detect_rooms_from_walls(
+    walls: list[Wall],
+    min_area: float = 1.0,
+    floor_boundary: Optional[Polygon] = None,
+) -> list[Polygon]:
+    """Detect rooms as enclosed regions between walls.
+    
+    Algorithm:
+    1. Build centerline graph from all walls
+    2. Use Shapely polygonize() to find all enclosed faces
+    3. Filter by minimum area
+    
+    For thick walls, we also try the subtraction approach:
+    1. Union all wall polygons
+    2. Subtract from floor boundary (or convex hull)
+    3. Remaining polygons = rooms
+    
+    Returns both approaches merged and deduplicated.
+    """
+    rooms: list[Polygon] = []
+
+    # Approach 1: Centerline polygonize
+    centerlines = [centerline_to_linestring(w) for w in walls]
+    merged_lines = unary_union(centerlines)
+
+    # Ensure we have a collection of lines for polygonize
+    if isinstance(merged_lines, LineString):
+        merged_lines = MultiLineString([merged_lines])
+
+    centerline_rooms = list(polygonize(merged_lines))
+    rooms.extend([r for r in centerline_rooms if r.area >= min_area])
+
+    # Approach 2: Wall subtraction (handles thick walls better)
+    if walls:
+        all_walls = wall_union(walls)
+        if floor_boundary is None:
+            # Use convex hull of all wall polygons + some margin
+            floor_boundary = all_walls.convex_hull.buffer(0.01)
+
+        floor_minus_walls = floor_boundary.difference(all_walls)
+        if not floor_minus_walls.is_valid:
+            floor_minus_walls = floor_minus_walls.buffer(0)
+
+        if isinstance(floor_minus_walls, MultiPolygon):
+            subtraction_rooms = [
+                g for g in floor_minus_walls.geoms if g.area >= min_area
+            ]
+        elif isinstance(floor_minus_walls, Polygon) and floor_minus_walls.area >= min_area:
+            subtraction_rooms = [floor_minus_walls]
+        else:
+            subtraction_rooms = []
+
+        # If centerline approach found rooms, use those (more precise topology)
+        # If not, fall back to subtraction approach
+        if not rooms and subtraction_rooms:
+            rooms = subtraction_rooms
+        elif subtraction_rooms and not rooms:
+            rooms = subtraction_rooms
+
+    return rooms
+
+
+def assign_room_wall_faces(
+    room_polygon: Polygon,
+    walls: list[Wall],
+    tolerance: float = 0.05,
+) -> list[WallFaceRef]:
+    """Determine which wall faces form a room's boundary.
+    
+    For each wall, check if the left or right offset line is adjacent to
+    (within tolerance of) the room polygon boundary.
+    """
+    boundary_refs: list[WallFaceRef] = []
+    room_boundary = room_polygon.boundary
+
+    for wall in walls:
+        left_line, right_line = wall_to_left_right_lines(wall)
+
+        # Check if left face touches the room
+        if isinstance(left_line, LineString) and left_line.length > 0:
+            dist_left = left_line.distance(room_boundary)
+            if dist_left < tolerance:
+                boundary_refs.append(
+                    WallFaceRef(wall_id=wall.id, side=WallSide.LEFT)
+                )
+                continue  # a wall typically only has one face per room
+
+        # Check if right face touches the room
+        if isinstance(right_line, LineString) and right_line.length > 0:
+            dist_right = right_line.distance(room_boundary)
+            if dist_right < tolerance:
+                boundary_refs.append(
+                    WallFaceRef(wall_id=wall.id, side=WallSide.RIGHT)
+                )
+
+    return boundary_refs
+
+
+def build_rooms(
+    walls: list[Wall],
+    min_area: float = 1.0,
+    floor_boundary: Optional[Polygon] = None,
+) -> tuple[list[Room], list[Polygon]]:
+    """Full room detection pipeline: find room polygons, assign wall faces.
+    
+    Returns (rooms, room_polygons) — the Room objects and their corresponding Shapely polygons.
+    """
+    room_polygons = detect_rooms_from_walls(
+        walls, min_area=min_area, floor_boundary=floor_boundary
+    )
+
+    rooms: list[Room] = []
+    for i, rpoly in enumerate(room_polygons):
+        wall_faces = assign_room_wall_faces(rpoly, walls)
+        room = Room(
+            id=f"r{i + 1}",
+            label=RoomLabel.UNKNOWN,
+            boundary=wall_faces if len(wall_faces) >= 2 else wall_faces,
+            area=round(rpoly.area, 2),
+        )
+        rooms.append(room)
+
+    return rooms, room_polygons
+
+
+# ──────────────────────────────────────────────
+# Floor plan assembly
+# ──────────────────────────────────────────────
+
+def compute_floor_plan_geometry(floorplan: FloorPlan) -> dict:
+    """Compute all derived geometry for a floor plan.
+    
+    Returns dict with:
+    - wall_polygons: {wall_id: Polygon}
+    - wall_polygons_with_openings: {wall_id: Polygon|MultiPolygon}
+    - room_polygons: [Polygon]
+    - wall_union: Polygon|MultiPolygon
+    """
+    wall_polys = {}
+    wall_polys_openings = {}
+
+    for wall in floorplan.walls:
+        wall_polys[wall.id] = wall_to_polygon(wall)
+        if wall.openings:
+            wall_polys_openings[wall.id] = wall_polygon_with_openings(wall)
+        else:
+            wall_polys_openings[wall.id] = wall_polys[wall.id]
+
+    all_walls = unary_union(list(wall_polys.values())) if wall_polys else Polygon()
+
+    room_polygons = detect_rooms_from_walls(floorplan.walls)
+
+    return {
+        "wall_polygons": wall_polys,
+        "wall_polygons_with_openings": wall_polys_openings,
+        "room_polygons": room_polygons,
+        "wall_union": all_walls,
+    }