73 lines
3.0 KiB
Python
73 lines
3.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
from src.distance_estimator import DistanceEstimator
|
|
from src.person_detector import PersonDetector
|
|
|
|
def test_distance_estimator():
|
|
# 初始化距离估算器（使用实际相机参数）
|
|
estimator = DistanceEstimator(
|
|
focal_length_mm=35,
|
|
sensor_width_mm=23.5,
|
|
sensor_height_mm=15.6,
|
|
image_width_pixels=1920,
|
|
image_height_pixels=1080,
|
|
camera_height_mm=1700,
|
|
camera_tilt_angle=15
|
|
)
|
|
|
|
print("测试距离估算...")
|
|
|
|
# 读取测试图像
|
|
test_image = cv2.imread("tests/test_data/test_image.jpg")
|
|
assert test_image is not None, "无法加载测试图像"
|
|
|
|
# 初始化检测器
|
|
detector = PersonDetector("yolov8n.pt", conf_threshold=0.3)
|
|
persons, vis_path = detector.detect(test_image, save_visualization=True)
|
|
|
|
print("len(person)", len(persons))
|
|
|
|
for person in persons:
|
|
assert 'bbox' in person, "检测结果缺少bbox"
|
|
assert 'confidence' in person, "检测结果缺少confidence"
|
|
assert len(person['bbox']) == 4, "bbox格式不正确"
|
|
assert 0 <= person['confidence'] <= 1, "confidence值范围不正确"
|
|
|
|
|
|
|
|
# 测试边界框方法
|
|
bbox = person['bbox'] # 示例边界框
|
|
bbox_distance = estimator.estimate_distance_from_bbox(bbox[3] - bbox[1])
|
|
print(f"边界框方法估算距离: {bbox_distance:.2f}mm")
|
|
assert bbox_distance > 0, "距离估算结果应该大于0"
|
|
|
|
# 测试关键点方法
|
|
keypoint_distance, keypoint_conf = estimator.estimate_distance_from_keypoints(test_image)
|
|
if keypoint_distance is not None:
|
|
print(f"关键点方法估算距离: {keypoint_distance:.2f}mm, 置信度: {keypoint_conf:.2f}")
|
|
assert keypoint_distance > 0, "距离估算结果应该大于0"
|
|
assert 0 <= keypoint_conf <= 1, "置信度应该在0-1之间"
|
|
|
|
# 测试地平面投影方法
|
|
ground_distance = estimator.estimate_distance_from_ground_plane(bbox[3])
|
|
print(f"地平面投影方法估算距离: {ground_distance:.2f}mm")
|
|
assert ground_distance > 0, "距离估算结果应该大于0"
|
|
|
|
# 测试综合估计方法
|
|
final_distance, final_confidence = estimator.estimate_distance(test_image, bbox)
|
|
print(f"综合估计距离: {final_distance:.2f}mm, 置信度: {final_confidence:.2f}")
|
|
assert final_distance > 0, "距离估算结果应该大于0"
|
|
assert 0 <= final_confidence <= 1, "置信度应该在0-1之间"
|
|
|
|
# 测试可视化
|
|
vis_image = estimator.visualize_estimation(test_image, bbox, final_distance, final_confidence)
|
|
assert vis_image is not None, "可视化结果不应为空"
|
|
assert vis_image.shape == test_image.shape, "可视化图像尺寸应与原图相同"
|
|
|
|
# 保存可视化结果
|
|
cv2.imwrite("tests/test_data/distance_estimation_result.jpg", vis_image)
|
|
|
|
print("距离估算测试完成!")
|
|
|
|
if __name__ == "__main__":
|
|
test_distance_estimator() |