157 lines
6.5 KiB
Python
157 lines
6.5 KiB
Python
import os
|
|
import numpy as np
|
|
import cv2
|
|
from matplotlib.colors import LinearSegmentedColormap
|
|
|
|
|
|
class FemtoBoltViewer:
|
|
def __init__(self, depth_min=900, depth_max=1300):
|
|
self.depth_min = depth_min
|
|
self.depth_max = depth_max
|
|
|
|
# 自定义彩虹色 colormap
|
|
colors = ['fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue',
|
|
'fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue',
|
|
'fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue',
|
|
'fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue']
|
|
self.custom_cmap = LinearSegmentedColormap.from_list("custom_cmap", colors)
|
|
|
|
# SDK 设备句柄和配置
|
|
self.device_handle = None
|
|
self.pykinect = None
|
|
self.config = None
|
|
|
|
# 缓存数组
|
|
self.background = None
|
|
self.output_buffer = None
|
|
self._depth_filtered = None # 用于复用深度图过滤结果
|
|
self._blur_buffer = None # 用于复用高斯模糊结果
|
|
|
|
# OpenCV 窗口
|
|
cv2.namedWindow("Depth CV", cv2.WINDOW_NORMAL)
|
|
|
|
def _load_sdk(self):
|
|
try:
|
|
import pykinect_azure as pykinect
|
|
self.pykinect = pykinect
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
dll_path = os.path.join(base_dir, "..", "dll", "femtobolt", "k4a.dll")
|
|
self.pykinect.initialize_libraries(track_body=False, module_k4a_path=dll_path)
|
|
return True
|
|
except Exception as e:
|
|
print(f"加载 SDK 失败: {e}")
|
|
return False
|
|
|
|
def _configure_device(self):
|
|
self.config = self.pykinect.default_configuration
|
|
self.config.depth_mode = self.pykinect.K4A_DEPTH_MODE_NFOV_UNBINNED
|
|
self.config.camera_fps = self.pykinect.K4A_FRAMES_PER_SECOND_15
|
|
self.config.synchronized_images_only = False
|
|
self.device_handle = self.pykinect.start_device(config=self.config)
|
|
|
|
def _generate_contour_image(self, depth):
|
|
"""改进版 OpenCV 等高线渲染，梯度平滑、局部对比增强"""
|
|
try:
|
|
# 初始化 depth_filtered 缓冲区
|
|
if self._depth_filtered is None or self._depth_filtered.shape != depth.shape:
|
|
self._depth_filtered = np.zeros_like(depth, dtype=np.uint16)
|
|
|
|
np.copyto(self._depth_filtered, depth) # 直接覆盖，不生成新数组
|
|
depth_filtered = self._depth_filtered
|
|
depth_filtered[depth_filtered > self.depth_max] = 0
|
|
depth_filtered[depth_filtered < self.depth_min] = 0
|
|
height, width = depth_filtered.shape
|
|
|
|
# 背景缓存
|
|
if self.background is None or self.background.shape[:2] != (height, width):
|
|
background_gray = int(0.5 * 255 * 0.3 + 255 * (1 - 0.3))
|
|
self.background = np.ones((height, width, 3), dtype=np.uint8) * background_gray
|
|
grid_spacing = max(height // 20, width // 20, 10)
|
|
for x in range(0, width, grid_spacing):
|
|
cv2.line(self.background, (x, 0), (x, height-1), (255, 255, 255), 1)
|
|
for y in range(0, height, grid_spacing):
|
|
cv2.line(self.background, (0, y), (width-1, y), (255, 255, 255), 1)
|
|
|
|
# 初始化输出缓存和模糊缓存
|
|
self.output_buffer = np.empty_like(self.background)
|
|
self._blur_buffer = np.empty_like(self.background)
|
|
|
|
# 复用输出缓存，避免 copy()
|
|
np.copyto(self.output_buffer, self.background)
|
|
output = self.output_buffer
|
|
valid_mask = depth_filtered > 0
|
|
|
|
if np.any(valid_mask):
|
|
# 连续归一化深度值
|
|
norm_depth = np.zeros_like(depth_filtered, dtype=np.float32)
|
|
norm_depth[valid_mask] = (depth_filtered[valid_mask] - self.depth_min) / (self.depth_max - self.depth_min)
|
|
norm_depth = np.clip(norm_depth, 0, 1) ** 0.8 # Gamma增强
|
|
|
|
# 使用 colormap 映射
|
|
cmap_colors = (self.custom_cmap(norm_depth)[..., :3] * 255).astype(np.uint8)
|
|
output[valid_mask] = cmap_colors[valid_mask]
|
|
|
|
# Sobel 边界检测 + cv2.magnitude 替换 np.hypot
|
|
depth_uint8 = (norm_depth * 255).astype(np.uint8)
|
|
gx = cv2.Sobel(depth_uint8, cv2.CV_32F, 1, 0, ksize=3)
|
|
gy = cv2.Sobel(depth_uint8, cv2.CV_32F, 0, 1, ksize=3)
|
|
grad_mag = cv2.magnitude(gx, gy)
|
|
grad_mag = grad_mag.astype(np.uint8)
|
|
|
|
# 自适应局部对比度增强（向量化）
|
|
edge_mask = grad_mag > 30
|
|
output[edge_mask] = np.clip(output[edge_mask].astype(np.float32) * 1.5, 0, 255).astype(np.uint8)
|
|
|
|
# 高斯平滑，复用 dst 缓冲区
|
|
cv2.GaussianBlur(output, (3, 3), 0.3, dst=self._blur_buffer)
|
|
return self._blur_buffer
|
|
|
|
except Exception as e:
|
|
print(f"等高线渲染失败: {e}")
|
|
return None
|
|
|
|
def run(self):
|
|
if not self._load_sdk():
|
|
print("SDK 加载失败，程序退出")
|
|
return
|
|
|
|
self._configure_device()
|
|
print("FemtoBolt 深度相机启动成功，按 Ctrl+C 或 ESC 退出", self.config)
|
|
|
|
try:
|
|
while True:
|
|
capture = self.device_handle.update()
|
|
if capture is None:
|
|
continue
|
|
ret, depth_image = capture.get_depth_image()
|
|
if not ret or depth_image is None:
|
|
continue
|
|
|
|
final_img = self._generate_contour_image(depth_image)
|
|
if final_img is not None:
|
|
# 推迟裁剪到显示阶段
|
|
h, w = final_img.shape[:2]
|
|
target_width = h // 2
|
|
if w > target_width:
|
|
left = (w - target_width) // 2
|
|
right = left + target_width
|
|
cv2.imshow("Depth CV", final_img[:, left:right])
|
|
else:
|
|
cv2.imshow("Depth CV", final_img)
|
|
|
|
if cv2.waitKey(1) & 0xFF == 27:
|
|
break
|
|
|
|
except KeyboardInterrupt:
|
|
print("检测到退出信号，结束程序")
|
|
finally:
|
|
if self.device_handle:
|
|
self.device_handle.stop()
|
|
self.device_handle.close()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
viewer = FemtoBoltViewer(depth_min=500, depth_max=700)
|
|
viewer.run()
|