pyqtgraph实现matplotlib 2D压力热力图转3D热力图用于姿态估计

核心实现逻辑 #

原有2D热力图基于7（长）4（宽）个传感节点拼接而成，每个节点包含66的传感子阵列，整床总有效传感分辨率为42*24。3D热力图的核心是将压力值映射为三维空间的Z轴高度，同时保持和原2D图一致的jet色阶映射规则，最终通过pyqtgraph的OpenGL渲染组件实现可交互的3D效果。

实现步骤 #

• 数据读取与校验：完全复用原有JSON读取、节点完整性校验逻辑
• 整床矩阵拼接：将分散的74个66子阵列，按物理排布位置拼接为完整的二维压力矩阵，替代原有seaborn分块子图拼接的绘图逻辑
• 3D视图初始化：创建pyqtgraph的GL 3D视图组件，配置初始视角、背景色、交互规则
• 坐标与颜色映射：生成床垫平面对应的X/Y网格坐标，将压力值作为Z轴高度，把140~180区间的压力值线性映射到jet色卡，和原2D图的颜色规则保持一致
• 曲面渲染：加载3D曲面图元，关闭插值平滑保持传感单元的块状显示效果，添加基础光照提升立体感

注意：压力值裁剪、色卡映射的阈值必须和原2D代码保持140~180的范围，否则颜色对应的压力区间会出现偏差，影响姿态估计的视觉判断。

完整实现代码 #

import json
import numpy as np
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph.opengl as gl

# 基础配置参数（和原2D代码保持一致）
length = 7          # 床垫长度方向节点数
width = 4           # 床垫宽度方向节点数
node_sub_size = 6   # 单个节点的子传感阵列尺寸
vmin = 140          # 压力显示下限
vmax = 180          # 压力显示上限
height_scale = 0.3  # 3D凸起高度缩放系数，数值越大高度差越明显

# ---------------------- 数据读取与预处理 ----------------------
fsr_list = []
with open('Session_16-09-21-01-49-27_modified.json', 'r') as f:
    data = json.load(f)

for j in range(1, length*width + 1):
    node_frames = []
    for item in data:
        if item['num'] == j:
            node_arr = np.array(item['fsr']).reshape((node_sub_size, node_sub_size))
            node_frames.append(node_arr)
    fsr_list.append(node_frames)

# 节点数据完整性校验
min_frame_len = len(min(fsr_list, key=lambda x: len(x)))
if min_frame_len == 0:
    raise RuntimeError("JSON文件缺失部分节点数据，无法生成整床热力图")

# 拼接整床压力矩阵（默认取第0帧，动态播放时可替换帧索引）
full_pressure = np.zeros((length*node_sub_size, width*node_sub_size))
for node_idx in range(length*width):
    row_block = node_idx // width
    col_block = node_idx % width
    node_data = fsr_list[node_idx][0]
    full_pressure[
        row_block*node_sub_size : (row_block+1)*node_sub_size,
        col_block*node_sub_size : (col_block+1)*node_sub_size
    ] = node_data

# 压力值裁剪到固定显示范围
full_pressure = np.clip(full_pressure, vmin, vmax)

# ---------------------- 3D视图渲染 ----------------------
app = QtGui.QApplication([])
pg.setConfigOption('background', 'white')
pg.setConfigOption('foreground', 'black')

# 创建3D视图窗口
view = gl.GLViewWidget()
view.setWindowTitle('床垫压力3D热力图')
view.setCameraPosition(distance=100, elevation=40, azimuth=225)  # 初始视角参数，可自行调整
view.show()

# 生成网格坐标
x_coord = np.arange(0, width*node_sub_size)
y_coord = np.arange(0, length*node_sub_size)
X, Y = np.meshgrid(x_coord, y_coord)
Z = full_pressure * height_scale  # Z轴高度按缩放系数调整

# 颜色映射（复用原2D图的jet色卡）
cmap = pg.colormap.get('jet')
norm_val = (full_pressure - vmin) / (vmax - vmin)
face_colors = cmap.map(norm_val.flatten(), alpha=True).reshape((*Z.shape, 4))

# 创建3D曲面对象
surf = gl.GLSurfacePlotItem(
    x=x_coord,
    y=y_coord,
    z=Z,
    colors=face_colors,
    smooth=False,  # 关闭平滑插值，保持传感单元块状效果
    shader='shaded'  # 启用光照着色提升立体感
)
# 平移曲面到视图中心
surf.translate(-width*node_sub_size/2, -length*node_sub_size/2, 0)
view.addItem(surf)

if __name__ == '__main__':
    app.exec_()

扩展调整说明 #

• 动态帧播放：新增QTimer定时器，定时读取对应帧的压力矩阵，调用surf.setData(z=new_z, colors=new_colors)即可实现动态压力序列播放，无需重建视图
• 效果调整：修改height_scale参数可调整3D凸起的高度差，修改setCameraPosition参数可调整初始观察角度，所有视角参数均支持鼠标拖拽交互实时调整
• 色阶条添加：可在窗口侧边布局中加入pg.ColorBarItem组件，绑定jet色卡和140~180的数值范围，实现和2D图一致的颜色刻度提示
• 坐标轴隐藏：如果不需要显示3D坐标轴，可移除默认加载的轴组件，保持和原2D图一致的无坐标轴纯净显示效果

内容的提问来源于stack exchange，提问作者Jaspreet Singh