71 lines
2.0 KiB
Python
71 lines
2.0 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,
|
|
)
|
|
|
|
|
|
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)
|