bluepulsar
451 lines
18 KiB
Python
451 lines
18 KiB
Python
import os
|
||
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" # Windows 下禁用硬件加速，避免兼容性问题
|
||
|
||
import cv2
|
||
import orisys
|
||
import numpy as np
|
||
from scipy.ndimage import gaussian_filter
|
||
|
||
def create_colorbar(min_val, max_val, height=400, width=60, colormap=cv2.COLORMAP_JET):
|
||
"""
|
||
创建颜色条图像，显示深度值的颜色映射尺度。
|
||
|
||
Args:
|
||
min_val: 最小值
|
||
max_val: 最大值
|
||
height: 颜色条高度（默认400）
|
||
width: 颜色条宽度（默认60）
|
||
colormap: OpenCV颜色映射（默认JET）
|
||
|
||
Returns:
|
||
np.ndarray: 包含颜色条和标签的BGR图像
|
||
"""
|
||
# 创建颜色条：从0到255的渐变（从底部到顶部）
|
||
gradient = np.linspace(255, 0, height, dtype=np.uint8).reshape(-1, 1)
|
||
gradient = np.tile(gradient, (1, width))
|
||
colorbar = cv2.applyColorMap(gradient, colormap)
|
||
|
||
# 添加文本标签
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
font_scale = 0.4
|
||
thickness = 1
|
||
color = (255, 255, 255) # 白色文字
|
||
|
||
# 在顶部显示最大值
|
||
max_text = f"{max_val:.1f}"
|
||
(text_width, text_height), _ = cv2.getTextSize(max_text, font, font_scale, thickness)
|
||
cv2.putText(colorbar, max_text, (width - text_width - 5, text_height + 5),
|
||
font, font_scale, color, thickness, cv2.LINE_AA)
|
||
|
||
# 在中间显示中值
|
||
mid_val = (min_val + max_val) / 2
|
||
mid_text = f"{mid_val:.1f}"
|
||
(text_width, text_height), _ = cv2.getTextSize(mid_text, font, font_scale, thickness)
|
||
cv2.putText(colorbar, mid_text, (width - text_width - 5, height // 2 + text_height // 2),
|
||
font, font_scale, color, thickness, cv2.LINE_AA)
|
||
|
||
# 在底部显示最小值
|
||
min_text = f"{min_val:.1f}"
|
||
(text_width, text_height), _ = cv2.getTextSize(min_text, font, font_scale, thickness)
|
||
cv2.putText(colorbar, min_text, (width - text_width - 5, height - 5),
|
||
font, font_scale, color, thickness, cv2.LINE_AA)
|
||
|
||
return colorbar
|
||
|
||
def add_colorbar_to_image(image, min_val, max_val, colormap=cv2.COLORMAP_JET):
|
||
"""
|
||
将颜色条添加到图像右侧。
|
||
|
||
Args:
|
||
image: 原始图像
|
||
min_val: 最小值
|
||
max_val: 最大值
|
||
colormap: OpenCV颜色映射（默认JET）
|
||
|
||
Returns:
|
||
np.ndarray: 添加了颜色条的图像
|
||
"""
|
||
h, w = image.shape[:2]
|
||
colorbar = create_colorbar(min_val, max_val, height=h, width=60, colormap=colormap)
|
||
|
||
# 水平拼接图像和颜色条
|
||
combined = np.hstack([image, colorbar])
|
||
return combined
|
||
|
||
def estimate_depth(foe=None, flow=None, init_depth=None, min_radial_flow=0.01, T_z=1.0,
|
||
expansion_mask=None, contraction_mask=None):
|
||
"""
|
||
基于FOE（扩张焦点）和光流场估计深度。
|
||
|
||
Args:
|
||
foe: FOE坐标，可以是：
|
||
- (x, y) 元组或列表
|
||
- None（如果未提供，返回零深度图）
|
||
flow: 光流场，形状为 (H, W, 2) 的 numpy 数组，包含 u 和 v 分量
|
||
init_depth: 初始深度值，可以是：
|
||
- 标量：用于统一缩放深度
|
||
- numpy数组：用于逐像素参考深度（会进行混合）
|
||
- None：不进行缩放
|
||
min_radial_flow: 最小径向流阈值，避免除零（默认0.01）
|
||
T_z: 平移速度参数，用于深度计算（默认1.0）
|
||
expansion_mask: 扩张区域掩码（可选），如果提供，只在掩码内计算深度
|
||
contraction_mask: 收缩区域掩码（可选），如果提供，只在掩码内计算深度
|
||
|
||
Returns:
|
||
np.ndarray: 深度图，形状为 (H, W)
|
||
|
||
Examples:
|
||
# 使用FOE坐标和vnormal估计深度
|
||
depth = estimate_depth(foe=(cx, cy), flow=vnormal, init_depth=30)
|
||
|
||
# 使用FOE列表和掩码
|
||
depth = estimate_depth(foe=[cx, cy], flow=vnormal, expansion_mask=expansion_mask, contraction_mask=contraction_mask)
|
||
"""
|
||
if flow is None:
|
||
raise ValueError("flow parameter is required")
|
||
|
||
H, W = flow.shape[:2]
|
||
u, v = flow[..., 0], flow[..., 1]
|
||
|
||
# 创建有效掩码：扩张和收缩区域的并集
|
||
if expansion_mask is not None or contraction_mask is not None:
|
||
if expansion_mask is not None and contraction_mask is not None:
|
||
valid_mask = expansion_mask | contraction_mask
|
||
elif expansion_mask is not None:
|
||
valid_mask = expansion_mask
|
||
else:
|
||
valid_mask = contraction_mask
|
||
else:
|
||
# 如果没有提供掩码，使用全图
|
||
valid_mask = np.ones((H, W), dtype=bool)
|
||
|
||
# 提取FOE坐标
|
||
if foe is None:
|
||
# 如果没有提供FOE，返回零深度图
|
||
return np.zeros((H, W), dtype=np.float32)
|
||
|
||
if isinstance(foe, (tuple, list, np.ndarray)):
|
||
if len(foe) >= 2:
|
||
fx, fy = float(foe[0]), float(foe[1])
|
||
else:
|
||
raise ValueError(f"FOE coordinates must have at least 2 elements, got {len(foe)}")
|
||
else:
|
||
raise ValueError(f"FOE must be a tuple, list, or array, got {type(foe)}")
|
||
|
||
# 确保FOE坐标在图像范围内
|
||
fx = np.clip(fx, 0, W - 1)
|
||
fy = np.clip(fy, 0, H - 1)
|
||
|
||
# 创建坐标网格
|
||
yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
|
||
|
||
# 计算从FOE到每个点的径向向量
|
||
rx = xx - fx
|
||
ry = yy - fy
|
||
r = np.sqrt(rx**2 + ry**2)
|
||
|
||
# 避免除零：在FOE点附近使用小值
|
||
r_safe = np.maximum(r, 1.0)
|
||
|
||
# 计算径向单位向量
|
||
radial_unit_x = rx / r_safe
|
||
radial_unit_y = ry / r_safe
|
||
|
||
# 计算径向流（光流在径向方向上的投影）
|
||
# 正值表示远离FOE，负值表示朝向FOE
|
||
radial_flow = u * radial_unit_x + v * radial_unit_y
|
||
|
||
# 创建掩码：FOE附近区域（半径内）
|
||
foe_radius = 10 # FOE附近的半径（像素）
|
||
mask_foe_region = r < foe_radius
|
||
|
||
# 处理FOE附近区域：使用邻域的中位数径向流（更稳健）
|
||
if np.any(mask_foe_region):
|
||
# 创建FOE区域的扩展掩码（用于获取邻域值）
|
||
mask_foe_extended = r < (foe_radius + 5)
|
||
if np.any(mask_foe_extended & ~mask_foe_region):
|
||
# 使用FOE区域外的邻域值的中位数
|
||
neighbor_flow = radial_flow[mask_foe_extended & ~mask_foe_region]
|
||
if len(neighbor_flow) > 0:
|
||
radial_flow[mask_foe_region] = np.median(neighbor_flow)
|
||
else:
|
||
radial_flow[mask_foe_region] = 0.0
|
||
else:
|
||
# 如果没有邻域值，使用全局中位数
|
||
if np.any(~mask_foe_region):
|
||
radial_flow[mask_foe_region] = np.median(radial_flow[~mask_foe_region])
|
||
else:
|
||
radial_flow[mask_foe_region] = 0.0
|
||
|
||
# 深度计算：Z = (T_z * r) / radial_flow
|
||
# 当径向流为正（远离FOE）时，深度为正
|
||
radial_flow_safe = np.where(
|
||
np.abs(radial_flow) < min_radial_flow,
|
||
np.sign(radial_flow) * min_radial_flow,
|
||
radial_flow
|
||
)
|
||
|
||
# 计算相对深度（仅在有效掩码内）
|
||
depth = np.zeros((H, W), dtype=np.float32)
|
||
|
||
# 只在有效掩码内计算深度
|
||
if np.any(valid_mask):
|
||
depth_valid = (T_z * r[valid_mask]) / np.maximum(np.abs(radial_flow_safe[valid_mask]), min_radial_flow)
|
||
depth_valid = np.abs(depth_valid) # 使用绝对值
|
||
depth[valid_mask] = depth_valid
|
||
|
||
# 对于掩码外的像素，设置为掩码内的最低深度值
|
||
if np.any(valid_mask):
|
||
min_depth_in_mask = np.min(depth[valid_mask])
|
||
depth[~valid_mask] = min_depth_in_mask
|
||
else:
|
||
# 如果没有有效掩码，使用初始深度或0
|
||
depth.fill(init_depth if init_depth is not None and np.isscalar(init_depth) else 0.0)
|
||
|
||
# 处理FOE附近区域：使用平滑后的值
|
||
if np.any(mask_foe_region):
|
||
depth_smoothed = gaussian_filter(depth, sigma=2.0)
|
||
depth[mask_foe_region] = depth_smoothed[mask_foe_region]
|
||
|
||
# 应用高斯滤波平滑整个深度图，减少噪声和伪影
|
||
depth = gaussian_filter(depth, sigma=1.5)
|
||
|
||
# 缩放深度（如果提供了初始深度，只使用掩码内的值）
|
||
if init_depth is not None:
|
||
if np.isscalar(init_depth):
|
||
# 统一缩放：将深度中位数缩放到init_depth（只使用掩码内的值）
|
||
if np.any(valid_mask):
|
||
valid_depth = depth[valid_mask & (depth > 0)]
|
||
if len(valid_depth) > 0:
|
||
median_depth = np.median(valid_depth)
|
||
if median_depth > 0:
|
||
scale = init_depth / median_depth
|
||
depth *= scale
|
||
elif isinstance(init_depth, np.ndarray):
|
||
# 逐像素参考深度：计算缩放因子并混合（只使用掩码内的值）
|
||
if init_depth.shape == depth.shape and np.any(valid_mask):
|
||
valid_ref = valid_mask & (init_depth > 0) & (depth > 0)
|
||
if np.any(valid_ref):
|
||
# 计算缩放因子
|
||
scale = np.median(init_depth[valid_ref] / depth[valid_ref])
|
||
depth *= scale
|
||
|
||
# 可选：与参考深度混合
|
||
alpha = 0.3 # 混合因子
|
||
depth = alpha * depth + (1 - alpha) * init_depth
|
||
|
||
# 缩放后，确保掩码外的像素保持为掩码内的最低深度值
|
||
if np.any(valid_mask):
|
||
min_depth_in_mask = np.min(depth[valid_mask])
|
||
depth[~valid_mask] = min_depth_in_mask
|
||
|
||
# 后处理：处理异常值
|
||
depth = np.nan_to_num(depth, nan=0.0, posinf=1000.0, neginf=0.0)
|
||
|
||
# 裁剪到合理范围（使用百分位数，但只在掩码内）
|
||
if np.any(valid_mask) and np.any(depth[valid_mask] > 0):
|
||
p1, p99 = np.percentile(depth[valid_mask & (depth > 0)], [1, 99])
|
||
depth = np.clip(depth, p1, p99)
|
||
|
||
# 最终平滑：对FOE附近区域再次应用平滑，确保伪影被消除
|
||
depth_final_smooth = gaussian_filter(depth, sigma=1.0)
|
||
|
||
# 重新计算FOE附近掩码
|
||
rx = xx - fx
|
||
ry = yy - fy
|
||
r = np.sqrt(rx**2 + ry**2)
|
||
mask_foe_region = r < foe_radius
|
||
|
||
# 在FOE附近区域使用更平滑的值（但只在有效掩码内）
|
||
if np.any(mask_foe_region & valid_mask):
|
||
depth[mask_foe_region & valid_mask] = depth_final_smooth[mask_foe_region & valid_mask]
|
||
|
||
# 确保掩码外的像素保持为掩码内的最低深度值
|
||
if np.any(valid_mask):
|
||
min_depth_in_mask = np.min(depth[valid_mask])
|
||
depth[~valid_mask] = min_depth_in_mask
|
||
|
||
return depth.astype(np.float32)
|
||
|
||
|
||
|
||
def main():
|
||
|
||
|
||
# 方式1: 从视频文件读取
|
||
# sensor = orisys.Sensor("./s1.mp4", isstitch=True, cuda=True, verbose=True)
|
||
|
||
# 方式2: 从摄像头读取（取消注释使用）
|
||
sensor = orisys.Sensor(0, isstitch=True, cuda=True, verbose=True)
|
||
|
||
print("\n按 'q' 退出程序，按 'r' 重置光流追踪器")
|
||
|
||
while True:
|
||
# 步骤 2.1: 获取并拼接图像
|
||
sensor.get_img()
|
||
|
||
# 步骤 2.2: 计算形变场
|
||
sensor.compute_deformation()
|
||
|
||
fps, flow,vnormal, depth = sensor.read_info(
|
||
sensor.info.FPS,
|
||
sensor.info.VRAW,
|
||
sensor.info.VNORMAL,
|
||
sensor.info.DEPTH
|
||
)
|
||
|
||
# depth = get_depth_improved(vnormal, fps)
|
||
|
||
|
||
arrows = orisys.util.draw_arrows(
|
||
sensor.img.copy(), # 使用 copy() 避免修改原图
|
||
vnormal,
|
||
threshold=5,
|
||
grid_spacing=10,
|
||
arrow_scale=1.0
|
||
)
|
||
cv2.imshow("形变场 Deformation", arrows)
|
||
|
||
# # 绘制深度图（热力图）
|
||
# # 深度值缩放：中心值 30，范围 ±10（即 20-40 映射到 0-255）
|
||
|
||
|
||
mag_img = orisys.util.draw_magnitude_map(
|
||
sensor.img.copy(),
|
||
vnormal,
|
||
threshold=5,
|
||
colormap=cv2.COLORMAP_HOT,
|
||
downsample_factor = 1
|
||
)
|
||
cv2.imshow("magnitude map", mag_img)
|
||
|
||
# 计算扩张焦点
|
||
# foe_result = orisys.util.compute_foe(vnormal,min_magnitude=1,expansion_threshold=0.01,downsample_factor=1)
|
||
foe_result = sensor.get_foe(multiple_foe=False, min_magnitude=1, expansion_threshold=0.01, downsample_factor=1)
|
||
# 在图像上绘制焦点和区域
|
||
image_with_foe = orisys.util.draw_foe(sensor.img.copy(), foe_result)
|
||
cv2.imshow("Focus of Expansion", image_with_foe)
|
||
|
||
divergence = foe_result['divergence_map']
|
||
|
||
# 计算散度统计信息
|
||
div_min = np.min(divergence)
|
||
div_max = np.max(divergence)
|
||
div_median = np.median(divergence)
|
||
print(f"散度统计: min={div_min:8.4f}, median={div_median:8.4f}, max={div_max:8.4f}")
|
||
|
||
# 归一化散度值到0-255范围（使用绝对值，然后归一化）
|
||
div_abs = np.abs(divergence)
|
||
if div_abs.max() > div_abs.min():
|
||
div_normalized = ((div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
|
||
else:
|
||
div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
|
||
|
||
# 应用颜色映射（使用JET：蓝色=低散度，红色=高散度）
|
||
divergence_colored = cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET)
|
||
|
||
# 添加颜色条显示散度尺度
|
||
div_display_min = div_abs.min()
|
||
div_display_max = div_abs.max()
|
||
divergence_with_colorbar = add_colorbar_to_image(divergence_colored, div_display_min, div_display_max, cv2.COLORMAP_JET)
|
||
cv2.imshow("divergence map", divergence_with_colorbar)
|
||
|
||
# 检查是否找到FOE
|
||
if 'foe_point' in foe_result and foe_result.get('is_foe_found', False):
|
||
(cx, cy) = foe_result['foe_point']
|
||
# 使用FOE坐标和vnormal估计深度，传入扩张和收缩掩码
|
||
Z_rel = estimate_depth(
|
||
foe=(cx, cy),
|
||
flow=vnormal,
|
||
init_depth=30,
|
||
T_z=1.0,
|
||
expansion_mask=foe_result.get('expansion_mask'),
|
||
contraction_mask=foe_result.get('contraction_mask')
|
||
)
|
||
|
||
# 计算深度统计信息（用于颜色映射缩放）
|
||
z_min = np.min(Z_rel)
|
||
z_max = np.max(Z_rel)
|
||
z_median = np.median(Z_rel)
|
||
|
||
# 打印深度统计信息
|
||
print(f"FPS: {fps:5.1f}, FOE: ({cx:3d}, {cy:3d})")
|
||
print(f" 深度统计 (用于颜色映射缩放): min={z_min:8.2f}, median={z_median:8.2f}, max={z_max:8.2f}")
|
||
|
||
# 智能颜色映射缩放：根据统计信息设置合理的显示范围
|
||
# 如果最大值远大于中位数，说明存在异常值，使用百分位数范围
|
||
if z_max > 5 * z_median:
|
||
# 存在异常值：使用百分位数范围，以median为中心
|
||
z_p1 = np.percentile(Z_rel, 1)
|
||
z_p99 = np.percentile(Z_rel, 99)
|
||
z_p5 = np.percentile(Z_rel, 5)
|
||
z_p95 = np.percentile(Z_rel, 95)
|
||
|
||
# 计算以median为中心的对称范围
|
||
# 使用5%和95%百分位数来确定主要数据范围
|
||
range_low = z_median - z_p5
|
||
range_high = z_p95 - z_median
|
||
range_symmetric = max(range_low, range_high)
|
||
|
||
# 设置显示范围：median ± 对称范围，但不小于1%和99%百分位数
|
||
z_display_min = max(z_p1, z_median - 2 * range_symmetric)
|
||
z_display_max = min(z_p99, z_median + 2 * range_symmetric)
|
||
|
||
print(f" 使用百分位数范围 (避免异常值): [{z_display_min:8.2f}, {z_display_max:8.2f}]")
|
||
else:
|
||
# 正常范围：使用min-max
|
||
z_display_min = z_min
|
||
z_display_max = z_max
|
||
print(f" 使用完整范围: [{z_display_min:8.2f}, {z_display_max:8.2f}]")
|
||
|
||
# 裁剪深度值到显示范围并归一化
|
||
depth_clipped = np.clip(Z_rel, z_display_min, z_display_max)
|
||
depth_normalized = (depth_clipped - z_display_min) / (z_display_max - z_display_min + 1e-6)
|
||
depth_normalized = np.clip(depth_normalized, 0, 1)
|
||
|
||
# 使用JET颜色映射显示深度图（蓝色=低深度，红色=高深度）
|
||
depth_colored = cv2.applyColorMap((depth_normalized * 255).astype(np.uint8), cv2.COLORMAP_JET)
|
||
|
||
# 添加颜色条显示深度尺度
|
||
depth_with_colorbar = add_colorbar_to_image(depth_colored, z_display_min, z_display_max, cv2.COLORMAP_JET)
|
||
cv2.imshow("Depth Map", depth_with_colorbar)
|
||
else:
|
||
# 如果没有找到FOE，显示提示信息
|
||
print("No FOE found, skipping depth estimation")
|
||
# 显示一个默认的深度图
|
||
Z_rel = np.full((sensor.img.shape[0], sensor.img.shape[1]), 30.0, dtype=np.float32)
|
||
depth_colored = cv2.applyColorMap((Z_rel * 8).astype(np.uint8), cv2.COLORMAP_HOT)
|
||
|
||
# 添加颜色条
|
||
depth_with_colorbar = add_colorbar_to_image(depth_colored, 0.0, 30.0, cv2.COLORMAP_HOT)
|
||
cv2.imshow("Depth Map", depth_with_colorbar)
|
||
|
||
# 计算多个扩张焦点
|
||
# foe_result = orisys.util.compute_multiple_foe(vnormal, min_contour_area=3000,min_magnitude=1,expansion_threshold=0.001,downsample_factor=10)
|
||
|
||
# # 在图像上绘制多个焦点和区域
|
||
# image_with_foe = orisys.util.draw_multiple_foe(sensor.img.copy(), foe_result)
|
||
# cv2.imshow("Multiple FOE", image_with_foe)
|
||
|
||
# =====================================================================
|
||
# 4. 键盘控制
|
||
# =====================================================================
|
||
key = cv2.waitKey(1) & 0xFF
|
||
if key == ord("q"):
|
||
print("\n退出程序...")
|
||
break
|
||
if key == ord("r"):
|
||
# 重置光流追踪器（当追踪漂移时使用）
|
||
sensor.reset()
|
||
print("光流追踪器已重置")
|
||
|
||
# =========================================================================
|
||
# 5. 清理资源
|
||
# =========================================================================
|
||
sensor.disconnect() # 释放视频源和相关资源
|
||
cv2.destroyAllWindows()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|