117 lines
4.3 KiB
Python
117 lines
4.3 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']
|
||
self.cmap = LinearSegmentedColormap.from_list("custom_cmap", colors)
|
||
|
||
# SDK 设备句柄和配置
|
||
self.device_handle = None
|
||
self.pykinect = None
|
||
self.config = None
|
||
|
||
# 缓存背景+网格图像（仅生成一次）
|
||
self.background = None
|
||
|
||
# OpenCV 窗口
|
||
cv2.namedWindow("Depth CV")
|
||
|
||
def _load_sdk(self):
|
||
"""加载并初始化 FemtoBolt SDK"""
|
||
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", "bin", "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):
|
||
"""配置 FemtoBolt 深度相机"""
|
||
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.config.color_resolution = 0
|
||
self.device_handle = self.pykinect.start_device(config=self.config)
|
||
|
||
def _get_color_image(self, depth_image):
|
||
"""将原始深度图转换为叠加背景网格后的 RGB 彩色图像"""
|
||
h, w = depth_image.shape
|
||
# 第一次调用时生成灰色背景和白色网格
|
||
if self.background is None:
|
||
self.background = np.full((h, w, 3), 128, dtype=np.uint8) # 灰色 (0.5 -> 128)
|
||
# 绘制网格线
|
||
for x in range(w):
|
||
cv2.line(self.background, (x, 0), (x, h-1), (255, 255, 255), 1)
|
||
for y in range(h):
|
||
cv2.line(self.background, (0, y), (w-1, y), (255, 255, 255), 1)
|
||
|
||
# 生成深度掩码，仅保留指定范围内的像素
|
||
mask_valid = (depth_image >= self.depth_min) & (depth_image <= self.depth_max)
|
||
depth_clipped = np.clip(depth_image, self.depth_min, self.depth_max)
|
||
normed = (depth_clipped.astype(np.float32) - self.depth_min) / (self.depth_max - self.depth_min)
|
||
|
||
# 反转映射，保证颜色方向与之前一致
|
||
normed = 1.0 - normed
|
||
|
||
# 应用自定义 colormap，将深度值映射到 RGB
|
||
rgba = self.cmap(normed)
|
||
rgb = (rgba[..., :3] * 255).astype(np.uint8)
|
||
|
||
# 叠加：在背景上覆盖彩色深度图（掩码处不覆盖，保留灰色背景+网格）
|
||
final_img = self.background.copy()
|
||
final_img[mask_valid] = rgb[mask_valid]
|
||
return final_img
|
||
|
||
def run(self):
|
||
if not self._load_sdk():
|
||
print("SDK 加载失败，程序退出")
|
||
return
|
||
|
||
self._configure_device()
|
||
print("FemtoBolt 深度相机启动成功，按 Ctrl+C 或 ESC 退出")
|
||
|
||
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._get_color_image(depth_image)
|
||
|
||
# OpenCV 显示
|
||
cv2.imshow("Depth CV", final_img)
|
||
# 按 ESC 键退出
|
||
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=900, depth_max=1100)
|
||
viewer.run()
|