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
@@ -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)