import numpy as np import pytest from aruco.ground_plane import ( unproject_depth_to_points, detect_floor_plane, compute_consensus_plane, compute_floor_correction, refine_ground_from_depth, create_ground_diagnostic_plot, save_diagnostic_plot, FloorPlane, FloorCorrection, GroundPlaneConfig, GroundPlaneMetrics, ) 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) def test_detect_floor_plane_perfect(): # Create points on a perfect plane: y = -1.5 (floor at -1.5m) # Normal should be [0, 1, 0] (pointing up) # Plane eq: 0*x + 1*y + 0*z + d = 0 => y + d = 0 => -1.5 + d = 0 => d = 1.5 # Generate grid of points x = np.linspace(-1, 1, 10) z = np.linspace(0, 5, 10) xx, zz = np.meshgrid(x, z) yy = np.full_like(xx, -1.5) points = np.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=1) # Add some noise to make it realistic but within threshold rng = np.random.default_rng(42) points += rng.normal(0, 0.001, points.shape) result = detect_floor_plane(points, distance_threshold=0.01, seed=42) assert result is not None assert isinstance(result, FloorPlane) normal = result.normal d = result.d inliers = result.num_inliers # Normal could be [0, 1, 0] or [0, -1, 0] depending on RANSAC # But we usually want it pointing "up" relative to camera or just consistent # Open3D segment_plane doesn't guarantee orientation # Check if it's vertical (y-axis aligned) assert abs(normal[1]) > 0.9 # Check distance # If normal is [0, 1, 0], d should be 1.5 # If normal is [0, -1, 0], d should be -1.5 if normal[1] > 0: assert abs(d - 1.5) < 0.01 else: assert abs(d + 1.5) < 0.01 assert inliers == 100 def test_detect_floor_plane_with_outliers(): # 100 inliers on floor y=-1.0 inliers = np.zeros((100, 3)) inliers[:, 0] = np.random.uniform(-1, 1, 100) inliers[:, 1] = -1.0 inliers[:, 2] = np.random.uniform(1, 5, 100) # 50 outliers (walls, noise) outliers = np.random.uniform(-2, 2, (50, 3)) outliers[:, 1] = np.random.uniform(-0.5, 1.0, 50) # Above floor points = np.vstack([inliers, outliers]) result = detect_floor_plane(points, distance_threshold=0.02, seed=42) assert result is not None assert abs(result.normal[1]) > 0.9 # Vertical normal assert result.num_inliers >= 100 # Should find all inliers def test_detect_floor_plane_insufficient_points(): points = np.array([[0, 0, 0], [1, 0, 0]]) # Only 2 points result = detect_floor_plane(points) assert result is None def test_detect_floor_plane_no_plane(): # Random cloud points = np.random.uniform(-1, 1, (100, 3)) # With high threshold it might find something, but with low threshold and random points... # Actually RANSAC almost always finds *something* in 3 points. # But let's test that it runs without crashing. result = detect_floor_plane(points, distance_threshold=0.001, seed=42) # It might return None if it can't find enough inliers for a model # Open3D segment_plane usually returns a model even if bad. # We'll check our wrapper behavior. pass def test_compute_consensus_plane_simple(): # Two identical planes planes = [ FloorPlane(normal=np.array([0, 1, 0]), d=1.5), FloorPlane(normal=np.array([0, 1, 0]), d=1.5), ] result = compute_consensus_plane(planes) np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6) assert abs(result.d - 1.5) < 1e-6 def test_compute_consensus_plane_weighted(): # Two planes, one with more weight # Plane 1: normal [0, 1, 0], d=1.0 # Plane 2: normal [0, 1, 0], d=2.0 # Weights: [1, 3] -> weighted avg d should be (1*1 + 3*2)/4 = 7/4 = 1.75 planes = [ FloorPlane(normal=np.array([0, 1, 0]), d=1.0), FloorPlane(normal=np.array([0, 1, 0]), d=2.0), ] weights = [1.0, 3.0] result = compute_consensus_plane(planes, weights) np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6) assert abs(result.d - 1.75) < 1e-6 def test_compute_consensus_plane_averaging_normals(): # Two planes with slightly different normals # n1 = [0, 1, 0] # n2 = [0.1, 0.995, 0] (approx) n1 = np.array([0, 1, 0], dtype=np.float64) n2 = np.array([0.1, 1.0, 0], dtype=np.float64) n2 /= np.linalg.norm(n2) planes = [FloorPlane(normal=n1, d=1.0), FloorPlane(normal=n2, d=1.0)] result = compute_consensus_plane(planes) # Expected normal is roughly average (normalized) avg_n = (n1 + n2) / 2.0 avg_d = 1.0 # (1.0 + 1.0) / 2.0 norm = np.linalg.norm(avg_n) expected_n = avg_n / norm expected_d = avg_d / norm np.testing.assert_allclose(result.normal, expected_n, atol=1e-6) assert abs(result.d - expected_d) < 1e-6 def test_compute_consensus_plane_empty(): with pytest.raises(ValueError): compute_consensus_plane([]) def test_compute_consensus_plane_flip_normals(): # If one normal is flipped, it should be flipped back to align with the majority/first # n1 = [0, 1, 0] # n2 = [0, -1, 0] # d1 = 1.0 # d2 = -1.0 (same plane, just flipped normal) planes = [ FloorPlane(normal=np.array([0, 1, 0]), d=1.0), FloorPlane(normal=np.array([0, -1, 0]), d=-1.0), ] result = compute_consensus_plane(planes) # Should align to first one (arbitrary choice, but consistent) np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6) assert abs(result.d - 1.0) < 1e-6 def test_detect_floor_plane_vertical_normal_check(): # Create points on a vertical wall (normal [1, 0, 0]) # Should be rejected by normal check in refine loop, but detect_floor_plane itself # just returns the plane. The filtering happens in refine_ground_from_depth. # So let's test that detect_floor_plane returns it correctly. # Wall at x=2.0 y = np.linspace(-1, 1, 10) z = np.linspace(0, 5, 10) yy, zz = np.meshgrid(y, z) xx = np.full_like(yy, 2.0) points = np.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=1) result = detect_floor_plane(points, distance_threshold=0.01, seed=42) assert result is not None # Normal should be roughly [1, 0, 0] assert abs(result.normal[0]) > 0.9 assert abs(result.normal[1]) < 0.1 def test_compute_consensus_plane_outlier_rejection(): # 3 planes: 2 consistent, 1 outlier p1 = FloorPlane(normal=np.array([0, 1, 0], dtype=np.float64), d=1.0) p2 = FloorPlane(normal=np.array([0, 1, 0], dtype=np.float64), d=1.05) # Outlier: different d p3 = FloorPlane(normal=np.array([0, 1, 0], dtype=np.float64), d=5.0) planes = [p1, p2, p3] # Should reject p3 and average p1, p2 result = compute_consensus_plane(planes) np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6) # Average of 1.0 and 1.05 is 1.025 assert abs(result.d - 1.025) < 0.01 def test_compute_consensus_plane_outlier_rejection_angle(): # 3 planes: 2 consistent, 1 outlier (tilted) p1 = FloorPlane(normal=np.array([0, 1, 0], dtype=np.float64), d=1.0) p2 = FloorPlane(normal=np.array([0, 1, 0], dtype=np.float64), d=1.0) # Outlier: tilted 45 deg norm = np.array([0, 1, 1], dtype=np.float64) norm = norm / np.linalg.norm(norm) p3 = FloorPlane(normal=norm, d=1.0) planes = [p1, p2, p3] result = compute_consensus_plane(planes) np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6) assert abs(result.d - 1.0) < 1e-6 def test_compute_floor_correction_identity(): # Current floor is already at target # Target y = 0.0 # Current plane: normal [0, 1, 0], d = 0.0 (y = 0) current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=0.0) result = compute_floor_correction(current_plane, target_floor_y=0.0) assert result.valid np.testing.assert_allclose(result.transform, np.eye(4), atol=1e-6) def test_compute_floor_correction_translation_only(): # Current floor is at y = -1.0 # Plane eq: y + d = 0 => -1 + d = 0 => d = 1.0 # Target y = 0.0 # We need to move everything UP by 1.0 (Ty = 1.0) current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=1.0) result = compute_floor_correction( current_plane, target_floor_y=0.0, max_translation_m=2.0 ) assert result.valid expected = np.eye(4) expected[1, 3] = 1.0 np.testing.assert_allclose(result.transform, expected, atol=1e-6) def test_compute_floor_correction_rotation_only(): # Current floor is tilted 45 deg around Z # Normal is [-0.707, 0.707, 0] # Target normal is [0, 1, 0] # We need to rotate -45 deg around Z to align normals angle = np.deg2rad(45) c, s = np.cos(angle), np.sin(angle) # Normal rotated by 45 deg around Z from [0, 1, 0] # Rz(45) @ [0, 1, 0] = [-s, c, 0] = [-0.707, 0.707, 0] normal = np.array([-s, c, 0]) d = 0.0 # Passes through origin current_plane = FloorPlane(normal=normal, d=d) result = compute_floor_correction( current_plane, target_floor_y=0.0, max_rotation_deg=90.0 ) assert result.valid T_corr = result.transform # Check rotation part # Should be Rz(-45) angle_corr = np.deg2rad(-45) cc, ss = np.cos(angle_corr), np.sin(angle_corr) expected_R = np.array([[cc, -ss, 0], [ss, cc, 0], [0, 0, 1]]) np.testing.assert_allclose(T_corr[:3, :3], expected_R, atol=1e-6) # Translation should be 0 since d=0 and we rotate around origin (roughly) assert np.linalg.norm(T_corr[:3, 3]) < 1e-6 def test_compute_floor_correction_bounds(): # Request huge translation # Current floor y = -10.0 (d=10.0) # Target y = 0.0 # Need Ty = 10.0 # Max trans = 0.1 current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=10.0) result = compute_floor_correction( current_plane, target_floor_y=0.0, max_translation_m=0.1 ) assert not result.valid assert "exceeds limit" in result.reason def test_compute_floor_correction_relative(): # Current floor: normal [0, 1, 0], d=1.0 (y=-1.0) # Target plane: normal [0, 1, 0], d=2.0 (y=-2.0) # We want to move current to target. # Shift = current_d - target_d = 1.0 - 2.0 = -1.0 # So we move DOWN by 1.0. # New y = -1.0 - 1.0 = -2.0. Correct. current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=1.0) target_plane = FloorPlane(normal=np.array([0, 1, 0]), d=2.0) result = compute_floor_correction( current_plane, target_plane=target_plane, max_translation_m=2.0 ) assert result.valid # Translation should be -1.0 along Y np.testing.assert_allclose(result.transform[1, 3], -1.0, atol=1e-6) def test_compute_floor_correction_relative_large_offset(): # Current floor: d=100.0 (y=-100.0) # Target plane: d=100.0 (y=-100.0) # Target Y (absolute) = 0.0 # If we used absolute correction, shift would be 100.0 -> fail. # With relative correction, shift is 0.0 -> success. current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=100.0) target_plane = FloorPlane(normal=np.array([0, 1, 0]), d=100.0) result = compute_floor_correction( current_plane, target_floor_y=0.0, target_plane=target_plane, max_translation_m=0.1, ) assert result.valid np.testing.assert_allclose(result.transform[:3, 3], 0.0, atol=1e-6) def test_refine_ground_from_depth_disabled(): config = GroundPlaneConfig(enabled=False) extrinsics = {"cam1": np.eye(4)} camera_data = {"cam1": {"depth": np.zeros((10, 10)), "K": np.eye(3)}} new_extrinsics, metrics = refine_ground_from_depth(camera_data, extrinsics, config) assert not metrics.success assert "disabled" in metrics.message assert new_extrinsics == extrinsics def test_refine_ground_from_depth_insufficient_cameras(): # Only 1 camera, need 2 config = GroundPlaneConfig(min_valid_cameras=2, min_inliers=10) # Create fake depth map that produces a plane # Plane at y=-1.0 width, height = 20, 20 K = np.eye(3) K[0, 2] = 10 K[1, 2] = 10 K[0, 0] = 20 K[1, 1] = 20 # Generate points on plane y=-1.0 # In camera frame (assuming cam at origin looking -Z), floor is at Y=-1.0 # But wait, standard camera frame is Y-down. # Let's assume world frame is Y-up. # If cam is at origin, and looking down -Z (OpenGL) or +Z (OpenCV). # Let's use identity extrinsics -> cam frame = world frame. # World frame Y-up. # So we want points with y=-1.0. # But unproject_depth gives points in camera frame. # If we want world y=-1.0, and T=I, then cam y=-1.0. # But unproject uses OpenCV convention: Y-down. # So y=-1.0 means 1m UP in camera frame. # Let's just make points that form A plane, doesn't matter which one, # as long as it's detected. # Let's make a flat plane at Z=2.0 (fronto-parallel) # This corresponds to normal [0, 0, 1] in camera frame. # With T=I, this is [0, 0, 1] in world frame. # This is NOT vertical (y-axis aligned). # So it gets rejected by our new normal_vertical_thresh check! # We need to make a plane that has normal roughly [0, 1, 0]. # Let's rotate the camera so that Z=2 plane becomes Y=-2 plane in world. # Rotate -90 deg around X. Rx_neg90 = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]) T_world_cam = np.eye(4) T_world_cam[:3, :3] = Rx_neg90 depth_map = np.full((height, width), 2.0, dtype=np.float32) # Need to ensure we have enough points for RANSAC # 20x20 = 400 points. # Stride default is 8. 20/8 = 2. 2x2 = 4 points. # RANSAC n=3. So 4 points is enough. # But min_inliers=10. 4 < 10. # So we need to reduce stride or increase size. config.stride = 1 camera_data = {"cam1": {"depth": depth_map, "K": K}} extrinsics = {"cam1": T_world_cam} new_extrinsics, metrics = refine_ground_from_depth(camera_data, extrinsics, config) assert not metrics.success assert "Found 1 valid planes" in metrics.message assert metrics.num_cameras_valid == 1 def test_refine_ground_from_depth_success(): # 2 cameras, both seeing floor at y=-1.0 # We want to correct it to y=0.0 config = GroundPlaneConfig( min_valid_cameras=2, min_inliers=10, target_y=0.0, max_translation_m=2.0, ransac_dist_thresh=0.05, ) width, height = 20, 20 K = np.eye(3) K[0, 2] = 10 K[1, 2] = 10 K[0, 0] = 20 K[1, 1] = 20 # Create points on plane y=-1.0 in WORLD frame # Cam 1 at origin. T_world_cam = I. # So points in cam 1 should be at y=-1.0. # OpenCV cam: Y-down. So y=-1.0 is UP. # Let's just use the fact that we transform points to world frame before detection. # So if we make depth map such that unprojected points + extrinsics -> plane y=-1.0. # Let's manually mock the detection to avoid complex depth map math # We can't easily mock internal functions without monkeypatching. # Instead, let's construct a depth map that corresponds to a plane. # Simplest: Camera looking down at floor. # Cam at (0, 2, 0) looking at (0, 0, 0). # World floor at y=0. # Cam floor distance = 2.0. # But here we want to simulate a MISALIGNED floor. # Say we think floor is at y=-1.0 (in our current world frame). # So we generate points at y=-1.0. # Let's try a simpler approach: # Create depth map for a plane Z=2.0 in camera frame. # Set extrinsics such that this plane becomes Y=-1.0 in world frame. # Plane Z=2.0 in cam frame: (x, y, 2). # We want R * (x, y, 2) + t = (X, -1, Z). # Let R = Rotation(-90 deg around X). # R = [[1, 0, 0], [0, 0, 1], [0, -1, 0]] # R * (x, y, 2) = (x, 2, -y). # We want Y_world = -1. # So 2 + ty = -1 => ty = -3. # So if we put cam at y=-3, rotated -90X. # Then Z=2 plane becomes Y=-1 plane. # Rotation -90 deg around X Rx_neg90 = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]) t = np.array([0, -3, 0]) T_world_cam = np.eye(4) T_world_cam[:3, :3] = Rx_neg90 T_world_cam[:3, 3] = t # Depth map: constant 2.0 depth_map = np.full((height, width), 2.0, dtype=np.float32) # Need to ensure we have enough points for RANSAC # 20x20 = 400 points. # Stride default is 8. 20/8 = 2. 2x2 = 4 points. # RANSAC n=3. So 4 points is enough. # But min_inliers=10. 4 < 10. # So we need to reduce stride or increase size. config.stride = 1 camera_data = { "cam1": {"depth": depth_map, "K": K}, "cam2": {"depth": depth_map, "K": K}, } extrinsics = { "cam1": T_world_cam, "cam2": T_world_cam, } new_extrinsics, metrics = refine_ground_from_depth(camera_data, extrinsics, config) assert metrics.success assert metrics.num_cameras_valid == 2 assert metrics.correction_applied # We started with floor at y=-1.0. Target is y=0.0. # So we expect translation of +1.0 in Y. # T_corr should have ty approx 1.0. # BUT wait, we changed the logic to be relative to consensus! # In this test, both cameras see floor at y=-1.0. # So consensus plane is at y=-1.0 (d=1.0). # Each camera sees floor at y=-1.0 (d=1.0). # Relative correction: shift = current_d - consensus_d = 1.0 - 1.0 = 0.0. # So NO correction is applied if we only align to consensus! # This confirms our change works as intended (aligns to consensus). # But the test expects alignment to target_y=0.0. # If we want to test that it aligns to consensus, we need to make them disagree. # Or we accept that if they agree, correction is 0. T_corr = metrics.camera_corrections["cam1"] # assert abs(T_corr[1, 3] - 1.0) < 0.1 # Old expectation assert abs(T_corr[1, 3]) < 0.1 # New expectation: 0 correction because they agree # Check new extrinsics # Should be unchanged T_new = new_extrinsics["cam1"] # assert abs(T_new[1, 3] - (-2.0)) < 0.1 # Old expectation assert abs(T_new[1, 3] - (-3.0)) < 0.1 # New expectation: unchanged (-3.0) # Verify per-camera corrections assert "cam1" in metrics.camera_corrections assert "cam2" in metrics.camera_corrections assert len(metrics.camera_corrections) == 2 def test_refine_ground_from_depth_partial_success(): # 2 cameras, but only 1 finds a plane # Should still succeed if min_valid_cameras=1 (or if we relax it) # But config default is 2. # Let's set min_valid_cameras=1 config = GroundPlaneConfig( min_valid_cameras=1, min_inliers=10, target_y=0.0, max_translation_m=2.0, ) width, height = 20, 20 K = np.eye(3) K[0, 2] = 10 K[1, 2] = 10 K[0, 0] = 20 K[1, 1] = 20 # Cam 1: Valid plane (same as success test) Rx_neg90 = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]) t = np.array([0, -3, 0]) T_world_cam = np.eye(4) T_world_cam[:3, :3] = Rx_neg90 T_world_cam[:3, 3] = t depth_map_valid = np.full((height, width), 2.0, dtype=np.float32) # Need to ensure we have enough points for RANSAC config.stride = 1 # Cam 2: Invalid depth (zeros) depth_map_invalid = np.zeros((height, width), dtype=np.float32) camera_data = { "cam1": {"depth": depth_map_valid, "K": K}, "cam2": {"depth": depth_map_invalid, "K": K}, } extrinsics = { "cam1": T_world_cam, "cam2": T_world_cam, } new_extrinsics, metrics = refine_ground_from_depth(camera_data, extrinsics, config) assert metrics.success assert metrics.num_cameras_valid == 1 assert metrics.correction_applied # Cam 1 should be corrected assert "cam1" in metrics.camera_corrections assert "cam1" not in metrics.skipped_cameras # Cam 2 should be skipped assert "cam2" not in metrics.camera_corrections assert "cam2" in metrics.skipped_cameras # Cam 2 extrinsics should be unchanged np.testing.assert_array_equal(new_extrinsics["cam2"], extrinsics["cam2"]) # Cam 1 extrinsics should be unchanged because it agrees with itself (consensus of 1) np.testing.assert_array_equal(new_extrinsics["cam1"], extrinsics["cam1"]) def test_create_ground_diagnostic_plot_smoke(): # Create minimal metrics and data metrics = GroundPlaneMetrics( success=True, consensus_plane=FloorPlane(normal=np.array([0, 1, 0]), d=1.0), ) camera_data = { "cam1": { "depth": np.full((10, 10), 2.0, dtype=np.float32), "K": np.eye(3), } } extrinsics_before = {"cam1": np.eye(4)} extrinsics_after = {"cam1": np.eye(4)} extrinsics_after["cam1"][1, 3] = 1.0 import plotly.graph_objects as go fig = create_ground_diagnostic_plot( metrics, camera_data, extrinsics_before, extrinsics_after ) assert isinstance(fig, go.Figure) # Check for some expected traces trace_names = [t.name for t in fig.data] assert any("World X" in name for name in trace_names if name) assert any("Consensus Plane" in name for name in trace_names if name) assert any("Points cam1" in name for name in trace_names if name) assert any("Cam cam1 (before)" in name for name in trace_names if name) assert any("Cam cam1 (after)" in name for name in trace_names if name) def test_save_diagnostic_plot_smoke(tmp_path): import plotly.graph_objects as go fig = go.Figure() plot_path = tmp_path / "diag.html" save_diagnostic_plot(fig, str(plot_path)) assert plot_path.exists() assert plot_path.stat().st_size > 0