📚 课程目录第4章姿态估计

4. 姿态估计

任务目标

在Jetson Nano上使用TensorRT加速的轻量级姿态估计模型（MoveNet Thunder），对USB摄像头实时视频流进行人体关键点检测，并在画面中绘制骨架与关键点。要求检测帧率不低于10 FPS（Jetson Nano 4GB版本实测参考），模型输入尺寸为256×256。

环境与依赖

Jetson Nano（JetPack 4.6+，已安装CUDA 10.2、cuDNN 8、TensorRT 8）
Python 3.6+，OpenCV（pip install opencv-python）
NumPy、pycuda（pip install pycuda）
MoveNet Thunder .tflite模型文件（从TensorFlow Hub下载，或使用下方转换脚本从TensorFlow模型转换）

实现方案

由于Jetson Nano无法直接运行TensorFlow Lite的GPU委托（XNNPACK不支持），本案例采用TensorRT引擎直接加载TFLite模型的替代路径：使用tf2trt将TFLite转换为TensorRT引擎，推理时通过pycuda进行张量传输。模型输出为17个关键点的(y, x, confidence)三元组，后处理时按置信度阈值过滤并绘制。

关键完整代码

import cv2
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import time

# ---------- TensorRT 引擎加载 ----------
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def load_engine(engine_path):
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

engine = load_engine('movenet_thunder.engine')
context = engine.create_execution_context()
h_input = cuda.pagelocked_empty(1 * 3 * 256 * 256, dtype=np.float32)
h_output = cuda.pagelocked_empty(1 * 17 * 3, dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
stream = cuda.Stream()

# ---------- 预处理 ----------
def preprocess(frame):
    img = cv2.resize(frame, (256, 256))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 127.5 - 1.0
    return img.transpose(2, 0, 1).flatten()

# ---------- 后处理 ----------
KEYPOINT_EDGES = [(0,1),(0,2),(1,3),(2,4),(0,5),(0,6),(5,7),(7,9),(6,8),(8,10),(5,11),(6,12),(11,13),(13,15),(12,14),(14,16)]
def postprocess(output, orig_shape, conf_thresh=0.3):
    h, w = orig_shape
    kps = output.reshape(17, 3)
    points = []
    for y, x, conf in kps:
        if conf > conf_thresh:
            points.append((int(x * w), int(y * h), float(conf)))
        else:
            points.append(None)
    return points

def draw_skeleton(frame, points):
    h, w = frame.shape[:2]
    # 绘制关键点
    for p in points:
        if p is not None:
            cv2.circle(frame, (p[0], p[1]), 4, (0, 255, 0), -1)
    # 绘制骨架连线
    for a, b in KEYPOINT_EDGES:
        if points[a] and points[b]:
            cv2.line(frame, (points[a][0], points[a][1]), (points[b][0], points[b][1]), (0, 0, 255), 2)
    return frame

# ---------- 主循环 ----------
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

while True:
    ret, frame = cap.read()
    if not ret:
        break
    orig_shape = frame.shape[:2]
    t0 = time.time()
    # 推理
    np.copyto(h_input, preprocess(frame))
    cuda.memcpy_htod_async(d_input, h_input, stream)
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
    # 后处理
    points = postprocess(h_output, orig_shape)
    frame = draw_skeleton(frame, points)
    fps = 1.0 / (time.time() - t0)
    cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
    cv2.imshow('Pose Estimation', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

TensorRT引擎生成命令

首次运行前需将TFLite模型转换为TensorRT引擎（在Jetson Nano上执行）：

# 安装tf2trt（需TensorFlow 2.x CPU版本）
pip install tensorflow==2.11.0
pip install git+https://github.com/amir-abdi/tf2trt.git

# 转换（输入为1x256x256x3，输出为1x17x3）
python -c "
from tf2trt import convert
convert('movenet_thunder.tflite', 'movenet_thunder.engine', 
        input_shapes={'input': (1, 256, 256, 3)})
"

验证方法与判定标准

运行测试：执行上述Python脚本，确保摄像头画面正常显示。
功能验证：站在摄像头前1.5～2米处，身体大部分区域可见。观察画面中是否在头部、肩膀、手肘、手腕、髋部、膝盖、脚踝处绘制绿色圆点，并以红色线段连接形成骨架。
性能验证：画面左上角FPS数值应稳定在10以上（Jetson Nano 4GB版本，使用Thunder模型int8量化）。若低于10，可改用MoveNet Lightning模型（输入192×192）或降低摄像头分辨率至480p。
鲁棒性检查：侧身站立、抬手、下蹲时，关键点应能跟随动作变化，偶尔丢失后能快速恢复。

常见问题

CUDA内存不足：关闭其他GUI程序，或改用Lightning模型（引擎文件更小）。
检测不到人体：确认人体完整入画，且与背景有足够对比度；适当降低conf_thresh至0.2。
FPS过低：确认使用的是TensorRT引擎而非TFLite直接推理；检查是否开启CPU频率模式（sudo nvpmodel -m 0）。