BodyBalanceEvaluation/backend/tests/My_TestFemtobolt.py

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, ListedColormap

# 设置matplotlib支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


class FemtoBoltDynamicViewer:
    def __init__(self, depth_min=900, depth_max=1300):
        self.depth_min = depth_min
        self.depth_max = depth_max

        # 使用display_x.py的原始颜色映射算法
        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.mcmap = LinearSegmentedColormap.from_list("custom_cmap", colors)

        # SDK 设备句柄和配置
        self.device_handle = None
        self.pykinect = None
        self.config = None

    def _load_sdk(self):
        """加载并初始化 FemtoBolt SDK"""
        import pykinect_azure as pykinect
        base_dir = os.path.dirname(os.path.abspath(__file__))
        dll_path = os.path.join(base_dir, "..", "dll", "femtobolt", "k4a.dll")
        self.pykinect = pykinect
        self.pykinect.initialize_libraries(track_body=False, module_k4a_path=dll_path)

    def _configure_device(self):
        """配置 FemtoBolt 深度相机"""
        self.config = self.pykinect.default_configuration
        self.config.depth_mode = self.pykinect.K4A_DEPTH_MODE_NFOV_2X2BINNED
        self.config.color_format = self.pykinect.K4A_IMAGE_FORMAT_COLOR_BGRA32
        self.config.color_resolution = self.pykinect.K4A_COLOR_RESOLUTION_720P
        self.config.synchronized_images_only = False
        self.config.color_resolution = 0
        self.device_handle = self.pykinect.start_device(config=self.config)



    def run(self):
        """运行实时深度数据可视化 - 融合display_x.py原始算法"""
        self._load_sdk()
        self._configure_device()

        plt.ion()  # 打开交互模式
        plt.figure(figsize=(7, 7))  # 使用display_x.py的图形设置
        print("FemtoBolt 深度相机启动成功，关闭窗口或 Ctrl+C 退出")
        print(f"深度范围: {self.depth_min} - {self.depth_max} mm")

        try:
            frame_count = 0
            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

                # 使用display_x.py的原始算法处理深度数据
                depth = depth_image.copy()
                
                # 深度数据过滤 (根据输入参数动态设置)
                depth[depth > self.depth_max] = 0
                depth[depth < self.depth_min] = 0
                
                # 裁剪感兴趣区域 (与display_x.py完全一致)
                # depth = depth[50:200, 50:210]

                # 背景图 (与display_x.py完全一致)
                background = np.ones_like(depth) * 0.5  # 设定灰色背景
                
                # 使用 np.ma.masked_equal() 来屏蔽深度图中的零值 (与display_x.py完全一致)
                depth = np.ma.masked_equal(depth, 0)

                # 绘制背景 (与display_x.py完全一致)
                plt.imshow(background, origin='lower', cmap='gray', alpha=0.3)
                
                # 绘制白色栅格线，并将其置于底层 (与display_x.py完全一致)
                plt.grid(True, which='both', axis='both', color='white', linestyle='-', linewidth=1, zorder=0)
                
                # 绘制等高线图并设置原点在左下角 (根据输入参数动态设置)
                # 通过设置 zorder 来控制它们的层级。例如，设置 zorder=2 或更大的值来确保它们位于栅格线之上。
                plt.contourf(depth, levels=50, cmap=self.mcmap, vmin=self.depth_min, vmax=self.depth_max, origin='upper', zorder=2)
                
                # 更新显示 (与display_x.py完全一致)
                plt.pause(0.1)  # 暂停0.1秒
                plt.draw()      # 重绘图像
                plt.clf()       # 清除当前图像
                
                frame_count += 1
                if frame_count % 30 == 0:  # 每30帧打印一次信息
                    print(f"已处理 {frame_count} 帧")

        except KeyboardInterrupt:
            print("\n检测到退出信号，结束程序")
        except Exception as e:
            print(f"运行时错误: {e}")
        finally:
            # 清理资源
            if self.device_handle:
                try:
                    if hasattr(self.device_handle, 'stop'):
                        self.device_handle.stop()
                    if hasattr(self.device_handle, 'close'):
                        self.device_handle.close()
                except Exception as e:
                    print(f"设备关闭时出现错误: {e}")
            plt.ioff()  # 关闭交互模式
            plt.close('all')
            print("程序已安全退出")

    def save_current_frame(self, filename="depth_frame.png"):
        """保存当前帧到文件"""
        try:
            plt.savefig(filename, dpi=150, bbox_inches='tight')
            print(f"当前帧已保存到: {filename}")
        except Exception as e:
            print(f"保存帧失败: {e}")


if __name__ == "__main__":
    # 创建查看器实例
    viewer = FemtoBoltDynamicViewer(depth_min=700, depth_max=1000)
    
    print("=" * 50)
    print("FemtoBolt 深度相机动态可视化测试")
    print("基于 display_x.py 算法的实时成像")
    print("=" * 50)
    
    # 运行可视化
    viewer.run()