545 lines
17 KiB
Python
545 lines
17 KiB
Python
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,
|
|
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_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_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)
|
|
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": np.eye(4)}
|
|
|
|
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.
|
|
T_corr = metrics.camera_corrections["cam1"]
|
|
assert abs(T_corr[1, 3] - 1.0) < 0.1 # Allow some slack for RANSAC noise
|
|
|
|
# Check new extrinsics
|
|
# New T = T_corr @ Old T
|
|
# Old T origin y = -3.
|
|
# New T origin y should be -3 + 1 = -2.
|
|
T_new = new_extrinsics["cam1"]
|
|
assert abs(T_new[1, 3] - (-2.0)) < 0.1
|
|
|
|
# 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,
|
|
)
|
|
|
|
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 changed
|
|
assert not np.array_equal(new_extrinsics["cam1"], extrinsics["cam1"])
|