feat: implement ground_plane.py with floor detection and alignment primitives

2026-02-09 07:05:35 +00:00
parent fb9a2b40b4
commit 1d3266ec60
3 changed files with 291 additions and 0 deletions
@@ -5,6 +5,7 @@
 {"id":"py_workspace-291","title":"Create camera pose comparison script","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-08T07:51:14.710189364Z","created_by":"crosstyan","updated_at":"2026-02-08T07:53:52.647760731Z","closed_at":"2026-02-08T07:53:52.647760731Z","close_reason":"Implemented compare_pose_sets.py script and verified with provided command."}
 {"id":"py_workspace-2c1","title":"Add manual ground-plane overlay to visualize_extrinsics.py","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T16:15:17.432846006Z","created_by":"crosstyan","updated_at":"2026-02-07T16:16:18.287496896Z","closed_at":"2026-02-07T16:16:18.287496896Z","close_reason":"Implemented ground-plane overlay with CLI options and updated README."}
 {"id":"py_workspace-49i","title":"Add explicit validation for 4x4 transformation matrices in compare_pose_sets.py","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-09T03:22:47.591167295Z","created_by":"crosstyan","updated_at":"2026-02-09T03:23:53.008806228Z","closed_at":"2026-02-09T03:23:53.008806228Z","close_reason":"Added explicit validation for 4x4 transformation matrices in parse_pose() with context-aware error messages. Verified with existing data."}
+{"id":"py_workspace-4o7","title":"Implement ground_plane.py for floor detection and alignment","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-09T06:58:07.905247984Z","created_by":"crosstyan","updated_at":"2026-02-09T07:04:51.276602825Z","closed_at":"2026-02-09T07:04:51.276602825Z","close_reason":"Implemented ground_plane.py with core primitives and tests"}
 {"id":"py_workspace-62y","title":"Fix depth pooling fallback threshold","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T08:12:12.046607198Z","created_by":"crosstyan","updated_at":"2026-02-07T08:13:12.98625698Z","closed_at":"2026-02-07T08:13:12.98625698Z","close_reason":"Updated fallback threshold to strict comparison"}
 {"id":"py_workspace-6m5","title":"Robust Optimizer Implementation","status":"closed","priority":0,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T05:22:45.183574374Z","created_by":"crosstyan","updated_at":"2026-02-07T05:22:53.151871639Z","closed_at":"2026-02-07T05:22:53.151871639Z","close_reason":"Implemented robust optimizer with least_squares and soft_l1 loss, updated tests"}
 {"id":"py_workspace-6sg","title":"Document marker parquet structure","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T02:48:08.95742431Z","created_by":"crosstyan","updated_at":"2026-02-07T02:49:35.897152691Z","closed_at":"2026-02-07T02:49:35.897152691Z","close_reason":"Documented parquet structure in aruco/markers/PARQUET_FORMAT.md"}
@@ -12,8 +13,10 @@
 {"id":"py_workspace-98p","title":"Integrate multi-frame depth pooling into calibrate_extrinsics.py","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T07:59:35.333468652Z","created_by":"crosstyan","updated_at":"2026-02-07T08:06:37.662956356Z","closed_at":"2026-02-07T08:06:37.662956356Z","close_reason":"Implemented multi-frame depth pooling and verified with tests"}
 {"id":"py_workspace-a85","title":"Add CLI option for ArUco dictionary in calibrate_extrinsics.py","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-06T10:13:41.896728814Z","created_by":"crosstyan","updated_at":"2026-02-07T07:29:52.290976525Z","closed_at":"2026-02-07T07:29:52.290976525Z","close_reason":"Implemented multi-frame depth pooling in calibrate_extrinsics.py"}
 {"id":"py_workspace-afh","title":"Inspect tmp_visualizer.html camera layout","notes":"Inspected tmp_visualizer.html. cam_0 is at (0,0,0). cam_1 is at (1,0,0). cam_2 is at (0, 0.5, 1.0). Axes are RGB=XYZ. Layout matches expected synthetic geometry.","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T15:40:04.162565539Z","created_by":"crosstyan","updated_at":"2026-02-07T15:42:10.721124074Z","closed_at":"2026-02-07T15:42:10.721124074Z","close_reason":"Inspection complete. Layout matches synthetic input."}
+{"id":"py_workspace-aif","title":"Update visualization conventions documentation","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-09T04:20:14.893831963Z","created_by":"crosstyan","updated_at":"2026-02-09T04:22:07.154821825Z","closed_at":"2026-02-09T04:22:07.154821825Z","close_reason":"Updated documentation with current policy checklist, metadata details, and known pitfalls"}
 {"id":"py_workspace-cg4","title":"Implement geometry-first auto-align heuristic","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T16:48:33.048250646Z","created_by":"crosstyan","updated_at":"2026-02-07T16:53:54.772815505Z","closed_at":"2026-02-07T16:53:54.772815505Z","close_reason":"Closed"}
 {"id":"py_workspace-cg9","title":"Implement core alignment utilities (Task 1)","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-06T10:40:36.296030875Z","created_by":"crosstyan","updated_at":"2026-02-06T10:40:46.196825039Z","closed_at":"2026-02-06T10:40:46.196825039Z","close_reason":"Implemented compute_face_normal, rotation_align_vectors, and apply_alignment_to_pose in aruco/alignment.py"}
+{"id":"py_workspace-e09","title":"Implement aruco/depth_save.py","status":"in_progress","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-09T06:58:01.987010195Z","created_by":"crosstyan","updated_at":"2026-02-09T06:58:09.311371064Z"}
 {"id":"py_workspace-ecz","title":"Update visualization conventions docs with alignment details","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-08T07:47:49.633647436Z","created_by":"crosstyan","updated_at":"2026-02-08T07:48:25.728323257Z","closed_at":"2026-02-08T07:48:25.728323257Z","close_reason":"Added alignment methodology section to docs"}
 {"id":"py_workspace-ee1","title":"Implement depth-mode argument resolution in calibrate_extrinsics.py","status":"closed","priority":2,"issue_type":"task","owner":"crosstyan@outlook.com","created_at":"2026-02-07T06:31:03.430147225Z","created_by":"crosstyan","updated_at":"2026-02-07T06:33:43.204825053Z","closed_at":"2026-02-07T06:33:43.204825053Z","close_reason":"Implemented depth-mode argument resolution logic and verified with multiple test cases."}
 {"id":"py_workspace-f23","title":"Add --origin-axes-scale option to visualize_extrinsics.py","status":"closed","priority":2,"issue_type":"feature","owner":"crosstyan@outlook.com","created_at":"2026-02-08T05:37:35.228917793Z","created_by":"crosstyan","updated_at":"2026-02-08T05:38:31.173898101Z","closed_at":"2026-02-08T05:38:31.173898101Z","close_reason":"Implemented --origin-axes-scale option and verified with rendering."}
@@ -0,0 +1,218 @@
+import numpy as np
+from typing import Optional, Tuple, List
+from jaxtyping import Float
+from typing import TYPE_CHECKING
+import open3d as o3d
+
+if TYPE_CHECKING:
+    Vec3 = Float[np.ndarray, "3"]
+    Mat44 = Float[np.ndarray, "4 4"]
+    PointsNC = Float[np.ndarray, "N 3"]
+else:
+    Vec3 = np.ndarray
+    Mat44 = np.ndarray
+    PointsNC = np.ndarray
+
+
+def unproject_depth_to_points(
+    depth_map: np.ndarray,
+    K: np.ndarray,
+    stride: int = 1,
+    depth_min: float = 0.1,
+    depth_max: float = 10.0,
+) -> PointsNC:
+    """
+    Unproject a depth map to a point cloud.
+    """
+    h, w = depth_map.shape
+    fx = K[0, 0]
+    fy = K[1, 1]
+    cx = K[0, 2]
+    cy = K[1, 2]
+
+    # Create meshgrid of pixel coordinates
+    # Use stride to reduce number of points
+    u_coords = np.arange(0, w, stride)
+    v_coords = np.arange(0, h, stride)
+    u, v = np.meshgrid(u_coords, v_coords)
+
+    # Sample depth map
+    z = depth_map[0:h:stride, 0:w:stride]
+
+    # Filter by depth bounds
+    valid_mask = (z > depth_min) & (z < depth_max) & np.isfinite(z)
+
+    # Apply mask
+    z_valid = z[valid_mask]
+    u_valid = u[valid_mask]
+    v_valid = v[valid_mask]
+
+    # Unproject
+    x_valid = (u_valid - cx) * z_valid / fx
+    y_valid = (v_valid - cy) * z_valid / fy
+
+    # Stack into (N, 3) array
+    points = np.stack((x_valid, y_valid, z_valid), axis=-1)
+
+    return points.astype(np.float64)
+
+
+def detect_floor_plane(
+    points: PointsNC,
+    distance_threshold: float = 0.02,
+    ransac_n: int = 3,
+    num_iterations: int = 1000,
+    seed: Optional[int] = None,
+) -> Tuple[Optional[Vec3], float, int]:
+    """
+    Detect the floor plane from a point cloud using RANSAC.
+    Returns (normal, d, num_inliers) where plane is normal.dot(p) + d = 0.
+    """
+    if points.shape[0] < ransac_n:
+        return None, 0.0, 0
+
+    # Convert to Open3D PointCloud
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+
+    # Set seed for determinism if provided
+    if seed is not None:
+        o3d.utility.random.seed(seed)
+
+    # Segment plane
+    plane_model, inliers = pcd.segment_plane(
+        distance_threshold=distance_threshold,
+        ransac_n=ransac_n,
+        num_iterations=num_iterations,
+    )
+
+    if not plane_model:
+        return None, 0.0, 0
+
+    # plane_model is [a, b, c, d]
+    a, b, c, d = plane_model
+    normal = np.array([a, b, c], dtype=np.float64)
+
+    # Normalize normal (Open3D usually returns normalized, but be safe)
+    norm = np.linalg.norm(normal)
+    if norm > 1e-6:
+        normal /= norm
+        d /= norm
+
+    return normal, d, len(inliers)
+
+
+def compute_consensus_plane(
+    planes: List[Tuple[Vec3, float]],
+    weights: Optional[List[float]] = None,
+) -> Tuple[Vec3, float]:
+    """
+    Compute a consensus plane from multiple plane detections.
+    """
+    if not planes:
+        raise ValueError("No planes provided for consensus.")
+
+    n_planes = len(planes)
+    if weights is None:
+        weights = [1.0] * n_planes
+
+    if len(weights) != n_planes:
+        raise ValueError(
+            f"Weights length {len(weights)} must match planes length {n_planes}"
+        )
+
+    # Use the first plane as reference for orientation
+    ref_normal = planes[0][0]
+
+    accum_normal = np.zeros(3, dtype=np.float64)
+    accum_d = 0.0
+    total_weight = 0.0
+
+    for i, (normal, d) in enumerate(planes):
+        w = weights[i]
+
+        # Check orientation against reference
+        if np.dot(normal, ref_normal) < 0:
+            # Flip normal and d to align with reference
+            normal = -normal
+            d = -d
+
+        accum_normal += normal * w
+        accum_d += d * w
+        total_weight += w
+
+    if total_weight <= 0:
+        raise ValueError("Total weight must be positive.")
+
+    avg_normal = accum_normal / total_weight
+    avg_d = accum_d / total_weight
+
+    # Re-normalize normal
+    norm = np.linalg.norm(avg_normal)
+    if norm > 1e-6:
+        avg_normal /= norm
+        # Scale d by 1/norm to maintain plane equation consistency
+        avg_d /= norm
+    else:
+        # Fallback (should be rare if inputs are valid)
+        avg_normal = np.array([0.0, 1.0, 0.0])
+        avg_d = 0.0
+
+    return avg_normal, float(avg_d)
+
+
+from .alignment import rotation_align_vectors
+
+
+def compute_floor_correction(
+    current_floor_plane: Tuple[Vec3, float],
+    target_floor_y: float = 0.0,
+    max_rotation_deg: float = 5.0,
+    max_translation_m: float = 0.1,
+) -> Optional[Mat44]:
+    """
+    Compute the correction transform to align the current floor plane to the target floor height.
+    Constrains correction to pitch/roll and vertical translation only.
+    """
+    current_normal, current_d = current_floor_plane
+
+    # Target normal is always [0, 1, 0] (Y-up)
+    target_normal = np.array([0.0, 1.0, 0.0])
+
+    # 1. Compute rotation to align normals
+    try:
+        R_align = rotation_align_vectors(current_normal, target_normal)
+    except ValueError:
+        return None
+
+    # Check rotation magnitude
+    # Angle of rotation is acos((trace(R) - 1) / 2)
+    trace = np.trace(R_align)
+    # Clip to avoid numerical errors outside [-1, 1]
+    cos_angle = np.clip((trace - 1) / 2, -1.0, 1.0)
+    angle_rad = np.arccos(cos_angle)
+    angle_deg = np.rad2deg(angle_rad)
+
+    if angle_deg > max_rotation_deg:
+        return None
+
+    # 2. Compute translation
+    # We want to move points such that the floor is at y = target_floor_y
+    # Plane equation: n . p + d = 0
+    # Current floor at y = -current_d (if n=[0,1,0])
+    # We want new y = target_floor_y
+    # So shift = target_floor_y - (-current_d) = target_floor_y + current_d
+
+    t_y = target_floor_y + current_d
+
+    # Check translation magnitude
+    if abs(t_y) > max_translation_m:
+        return None
+
+    # Construct T
+    T = np.eye(4)
+    T[:3, :3] = R_align
+    # Translation is applied in the rotated frame (aligned to target normal)
+    T[:3, 3] = target_normal * t_y
+
+    return T.astype(np.float64)
@@ -0,0 +1,70 @@
+import numpy as np
+import pytest
+from aruco.ground_plane import (
+    unproject_depth_to_points,
+    detect_floor_plane,
+    compute_consensus_plane,
+    compute_floor_correction,
+)
+
+
+def test_unproject_depth_to_points_simple():
+    # Simple 3x3 depth map
+    # K = identity for simplicity (fx=1, fy=1, cx=1, cy=1)
+    # Pixel (1, 1) is center.
+    # At (1, 1), u=1, v=1. x = (1-1)/1 = 0, y = (1-1)/1 = 0.
+    # If depth is Z, point is (0, 0, Z).
+
+    width, height = 3, 3
+    K = np.array([[1, 0, 1], [0, 1, 1], [0, 0, 1]], dtype=np.float64)
+    depth_map = np.zeros((height, width), dtype=np.float32)
+
+    # Center pixel
+    depth_map[1, 1] = 2.0
+
+    # Top-left pixel (0, 0)
+    # u=0, v=0. x = (0-1)/1 = -1. y = (0-1)/1 = -1.
+    # Point: (-1*Z, -1*Z, Z)
+    depth_map[0, 0] = 1.0
+
+    points = unproject_depth_to_points(depth_map, K, depth_min=0.1, depth_max=5.0)
+
+    # Should have 2 points (others are 0.0 which is < depth_min)
+    assert points.shape == (2, 3)
+
+    # Check center point
+    # We don't know order, so check if expected points exist
+    expected_center = np.array([0.0, 0.0, 2.0])
+    expected_tl = np.array([-1.0, -1.0, 1.0])
+
+    # Find matches
+    has_center = np.any(np.all(np.isclose(points, expected_center, atol=1e-5), axis=1))
+    has_tl = np.any(np.all(np.isclose(points, expected_tl, atol=1e-5), axis=1))
+
+    assert has_center
+    assert has_tl
+
+
+def test_unproject_depth_to_points_stride():
+    width, height = 10, 10
+    K = np.eye(3)
+    depth_map = np.ones((height, width), dtype=np.float32)
+
+    points = unproject_depth_to_points(depth_map, K, stride=2)
+
+    # 10x10 -> 5x5 = 25 points
+    assert points.shape == (25, 3)
+
+
+def test_unproject_depth_to_points_bounds():
+    width, height = 3, 3
+    K = np.eye(3)
+    depth_map = np.array(
+        [[0.05, 1.0, 11.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.float32
+    )
+
+    # 0.05 < 0.1 (min) -> excluded
+    # 11.0 > 10.0 (max) -> excluded
+    # 7 valid points
+    points = unproject_depth_to_points(depth_map, K, depth_min=0.1, depth_max=10.0)
+    assert points.shape == (7, 3)