基于QT开发SDPose-Wholebody的本地可视化工具
如果你正在寻找一个能精准识别人体133个关键点的姿态估计模型,SDPose-Wholebody绝对值得一试。它基于Stable Diffusion的视觉先验,在艺术风格、动画等“非正常”图像上表现尤其出色。不过,官方提供的Gradio界面虽然方便,但如果你想把它集成到自己的桌面应用里,或者想做一个更灵活、功能更强大的本地工具,该怎么办呢?
今天,我们就来聊聊怎么用QT框架,亲手打造一个SDPose-Wholebody的本地可视化工具。这个工具不仅能加载模型、识别图片和视频中的人体姿态,还能让你实时查看、编辑关键点,甚至把结果保存下来。整个过程就像搭积木,一步步来,你会发现其实没那么复杂。
1. 为什么选择QT?项目目标与准备工作
在开始敲代码之前,我们先想清楚两件事:为什么选QT,以及我们要做个什么样的工具。
QT是一个老牌且强大的跨平台C++图形用户界面库。用它来开发这个工具,主要有几个好处:
- 跨平台:写一次代码,就能在Windows、macOS、Linux上运行,省心。
- 性能好:C++底层,处理图像、视频这种数据密集型任务很流畅。
- 控件丰富:按钮、滑块、图像显示区域(QLabel/QGraphicsView)等应有尽有,组装界面快。
- 生态成熟:文档全,社区活跃,遇到问题容易找到解决方案。
我们的目标是构建一个具备以下核心功能的桌面应用:
- 模型加载:能够加载SDPose-Wholebody预训练模型。
- 图像/视频输入:支持选择本地图片或视频文件,也支持摄像头实时捕获。
- 姿态估计与可视化:对输入内容进行133点全身姿态估计,并用骨架图(连线)和关键点(圆圈)清晰地画出来。
- 交互与编辑(进阶):允许用户点击调整不满意的关键点位置。
- 结果导出:将带有关键点标注的图像或视频保存下来。
环境准备首先,确保你的开发环境已经就绪:
- Python 3.8+:这是必须的。
- PyQt5 或 PySide6:这是QT的Python绑定。我个人更喜欢PySide6,因为许可更友好。安装命令:
pip install PySide6 - SDPose-Wholebody 运行环境:你需要能成功运行官方的SDPose-OOD代码。这意味着要安装好PyTorch、MMPose、Diffusers等依赖。可以参考项目README,通常就是
pip install -r requirements.txt。 - OpenCV:用于图像/视频的读取、处理和显示。
pip install opencv-python
准备好这些,我们的“造车”车间就算布置好了。
2. 设计工具界面:用QT Designer快速布局
直接手写界面代码比较繁琐,我们可以先用QT Designer这个可视化工具来拖拽出界面。这里我设计了一个简单的布局,包含以下几个主要区域:
- 菜单栏/工具栏:放置“打开文件”、“开始/停止”、“保存结果”等动作。
- 左侧控制面板:
- 文件选择按钮(图片/视频/摄像头)。
- 模型加载状态显示。
- 一些参数调节滑块(比如关键点显示阈值、骨架线条粗细)。
- 运行控制按钮(开始推理、暂停、重置)。
- 中央图像显示区域:这是核心,用一个
QGraphicsView和QGraphicsScene来承载和显示图像、绘制骨架。QGraphicsView支持缩放、平移,查看细节更方便。 - 右侧信息面板:
- 显示当前检测到的人数。
- 列出选中人物的关键点坐标和置信度。
- 日志输出框。
用QT Designer画好界面后,会生成一个.ui文件。我们可以用pyside6-uic命令将它转换为Python代码,或者直接在程序里动态加载。为了结构清晰,我建议使用动态加载的方式。
3. 核心功能实现:连接界面与SDPose引擎
界面是“壳”,SDPose模型是“发动机”。现在我们要把它们连接起来。我们创建一个主窗口类,来统筹一切。
import sys import cv2 from pathlib import Path from PySide6.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox from PySide6.QtCore import Qt, QTimer, Signal, QThread, Slot from PySide6.QtGui import QImage, QPixmap, QPainter, QPen, QColor # 假设你的ui文件叫 main_window.ui from ui_main_window import Ui_MainWindow # 导入SDPose相关的推理模块 # 这里需要根据SDPose-OOD项目的实际结构来调整导入路径 import sys sys.path.append('/path/to/SDPose-OOD') # 添加项目路径 from inference import SDPoseInference # 这是一个假设的推理类,你需要根据官方代码封装一个 class PoseEstimationThread(QThread): """ 在一个单独的线程中进行姿态估计,避免界面卡顿 """ frame_processed = Signal(object, list) # 信号:发送处理后的图像和关键点列表 def __init__(self, inference_engine): super().__init__() self.inference_engine = inference_engine self.current_frame = None self.is_running = False def run(self): while self.is_running: if self.current_frame is not None: # 调用SDPose引擎进行推理 # 注意:这里需要你根据SDPose的实际推理函数来写 # 假设 inference_engine.predict(frame) 返回 (annotated_frame, keypoints_list) annotated_frame, keypoints = self.inference_engine.predict(self.current_frame) self.frame_processed.emit(annotated_frame, keypoints) self.msleep(30) # 控制一下处理频率 def update_frame(self, frame): self.current_frame = frame class MainWindow(QMainWindow): def __init__(self): super().__init__() self.ui = Ui_MainWindow() self.ui.setupUi(self) # 初始化变量 self.current_image_path = None self.cap = None # 视频捕获对象 self.timer = QTimer() # 用于视频帧定时器 self.inference_engine = None self.pose_thread = None self.keypoints = [] # 存储当前帧的关键点结果 self.colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0), (255, 255, 0)] # 不同人的骨架颜色 # 连接信号与槽 self.connect_signals_slots() # 初始化图形场景(用于绘制) self.scene = QGraphicsScene() self.ui.graphicsView.setScene(self.scene) def connect_signals_slots(self): """ 连接界面按钮和对应的功能函数 """ self.ui.actionOpenImage.triggered.connect(self.open_image) self.ui.actionOpenVideo.triggered.connect(self.open_video) self.ui.actionLoadModel.triggered.connect(self.load_model) self.ui.actionStartCamera.triggered.connect(self.start_camera) self.ui.btnStartInference.clicked.connect(self.start_inference) self.ui.btnStopInference.clicked.connect(self.stop_inference) self.ui.btnSaveResult.clicked.connect(self.save_result) self.timer.timeout.connect(self.update_video_frame) def load_model(self): """ 加载SDPose-Wholebody模型 """ try: # 初始化推理引擎,这里需要你根据SDPose-OOD的代码来实例化 # 例如,指定模型checkpoint路径,设备(cuda/cpu)等 checkpoint_path = './models/sdpose_wholebody.pth' # 你的模型路径 self.inference_engine = SDPoseInference(checkpoint_path, device='cuda:0') QMessageBox.information(self, "成功", "模型加载成功!") self.ui.statusbar.showMessage("模型已就绪") except Exception as e: QMessageBox.critical(self, "错误", f"模型加载失败: {e}") def open_image(self): """ 打开图片文件 """ file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg *.bmp)") if file_path: self.current_image_path = file_path frame = cv2.imread(file_path) self.display_frame(frame) # 停止任何正在进行的视频或摄像头 self.stop_video_source() def open_video(self): """ 打开视频文件 """ file_path, _ = QFileDialog.getOpenFileName(self, "选择视频", "", "Videos (*.mp4 *.avi *.mov *.mkv)") if file_path: self.stop_video_source() # 先停止之前的视频源 self.cap = cv2.VideoCapture(file_path) if not self.cap.isOpened(): QMessageBox.warning(self, "警告", "无法打开视频文件") return self.timer.start(30) # 约30ms一帧,33 FPS def start_camera(self): """ 启动摄像头 """ self.stop_video_source() self.cap = cv2.VideoCapture(0) # 0 表示默认摄像头 if not self.cap.isOpened(): QMessageBox.warning(self, "警告", "无法打开摄像头") return self.timer.start(30) def stop_video_source(self): """ 停止视频或摄像头捕获 """ if self.timer.isActive(): self.timer.stop() if self.cap: self.cap.release() self.cap = None def update_video_frame(self): """ 定时器槽函数:读取并显示视频帧 """ if self.cap and self.cap.isOpened(): ret, frame = self.cap.read() if ret: self.display_frame(frame) else: self.timer.stop() self.cap.release() def display_frame(self, frame): """ 将OpenCV图像显示在QT的GraphicsView中 """ # 转换颜色空间 BGR -> RGB rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) h, w, ch = rgb_image.shape bytes_per_line = ch * w # 创建QImage qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888) pixmap = QPixmap.fromImage(qt_image) # 清空场景并添加新的Pixmap self.scene.clear() self.scene.addPixmap(pixmap) self.ui.graphicsView.fitInView(self.scene.itemsBoundingRect(), Qt.KeepAspectRatio) # 存储当前帧,供推理线程使用(如果开启了推理) self.current_frame_for_inference = frame.copy() @Slot(object, list) def on_frame_processed(self, annotated_frame, keypoints): """ 接收处理后的帧和关键点,并更新显示 """ self.keypoints = keypoints # 将标注后的帧(BGR格式)显示出来 self.display_frame(annotated_frame) # 更新右侧信息面板,例如人数 self.ui.label_person_count.setText(f"检测到人数: {len(keypoints)}") def start_inference(self): """ 开始姿态估计推理 """ if self.inference_engine is None: QMessageBox.warning(self, "警告", "请先加载模型") return if self.pose_thread is None: self.pose_thread = PoseEstimationThread(self.inference_engine) self.pose_thread.frame_processed.connect(self.on_frame_processed) if not self.pose_thread.isRunning(): self.pose_thread.is_running = True self.pose_thread.start() # 启动一个定时器,不断将当前帧送给推理线程 self.inference_timer = QTimer() self.inference_timer.timeout.connect(self.feed_frame_to_thread) self.inference_timer.start(50) # 20 FPS def feed_frame_to_thread(self): """ 定时将当前界面显示的帧发送给推理线程 """ if hasattr(self, 'current_frame_for_inference') and self.pose_thread: self.pose_thread.update_frame(self.current_frame_for_inference.copy()) def stop_inference(self): """ 停止推理 """ if self.pose_thread and self.pose_thread.isRunning(): self.pose_thread.is_running = False self.pose_thread.quit() self.pose_thread.wait() if hasattr(self, 'inference_timer'): self.inference_timer.stop() def save_result(self): """ 保存当前标注结果 """ if not self.keypoints: QMessageBox.information(self, "提示", "没有可保存的结果") return # 这里可以实现保存图片或视频的逻辑 # 例如,将当前scene渲染成图片保存 pass if __name__ == "__main__": app = QApplication(sys.argv) window = MainWindow() window.show() sys.exit(app.exec())这段代码搭建了基本的框架。核心是PoseEstimationThread类,它在一个独立的线程中运行SDPose推理,防止界面失去响应。主窗口类MainWindow负责界面交互、媒体文件处理,并通过信号槽机制与推理线程通信。
4. 关键点可视化与交互优化
把关键点画出来,并且画得好看、易懂,是可视化工具的灵魂。SDPose-Wholebody输出133个点,我们需要按照COCO-WholeBody的格式(身体17点、脸部68点、手部42点、脚部6点)来连接它们,形成骨架。
我们需要一个绘制函数,可以添加到display_frame之后,或者直接在on_frame_processed中,在annotated_frame上绘制。但更QT的方式是,在QGraphicsScene上叠加绘制,这样不会破坏原始图像,也便于做交互编辑。
def draw_skeleton_on_scene(self, keypoints_list): """ 在GraphicsScene上绘制骨架和关键点 """ # 先清除之前绘制的骨架图形项(如果有的话) for item in self.skeleton_items: self.scene.removeItem(item) self.skeleton_items.clear() if not keypoints_list: return # 定义关键点连接关系(COCO-WholeBody格式简化示例) # 这里需要你根据完整的133点连接表来完善 body_skeleton = [(0,1), (1,2), (2,3), ...] # 身体连接 face_skeleton = [...] # 脸部连接(通常不连线,或连成轮廓) hand_skeleton = [...] # 手部连接 foot_skeleton = [...] # 脚部连接 for person_idx, keypoints in enumerate(keypoints_list): color = self.colors[person_idx % len(self.colors)] qcolor = QColor(*color) # 绘制关键点(圆圈) for kp in keypoints: x, y, conf = kp # 假设每个关键点是(x, y, confidence)格式 if conf > 0.1: # 置信度阈值 ellipse = self.scene.addEllipse(x-3, y-3, 6, 6, QPen(qcolor), QBrush(qcolor)) self.skeleton_items.append(ellipse) # 绘制骨架连线 for start_idx, end_idx in body_skeleton: start_kp = keypoints[start_idx] end_kp = keypoints[end_idx] if start_kp[2] > 0.1 and end_kp[2] > 0.1: # 两点都可见 line = self.scene.addLine(start_kp[0], start_kp[1], end_kp[0], end_kp[1], QPen(qcolor, 2)) self.skeleton_items.append(line) # 同理绘制手、脚、脸的连线...交互编辑是一个进阶功能。思路是给QGraphicsScene上的每个关键点圆圈(QGraphicsEllipseItem)添加鼠标事件(点击、拖拽)。当用户拖拽一个点时,更新该点的坐标,并重新绘制与之相连的骨架线。这需要更细致的事件处理和状态管理。
5. 性能优化与实用技巧
当处理高分辨率图片或实时视频时,性能至关重要。这里有几个小建议:
- 图像缩放:SDPose模型输入分辨率是1024x768。如果原始图像很大,可以先在UI线程缩放到一个适中的预览尺寸(如800x600)进行显示,同时将原始尺寸或1024x768尺寸的图片送给推理线程。避免在UI线程进行大图的高频缩放。
- 推理线程管理:确保推理线程不会堆积帧。可以使用一个队列,并且只处理最新的帧,跳过中间积压的帧,这对于实时性要求高的场景(如摄像头)很有用。
- 绘制优化:在
draw_skeleton_on_scene中,批量更新图形项,而不是频繁地addEllipse和addLine。可以考虑自定义一个QGraphicsItem来代表一个人物的所有骨架,一次性绘制。 - 模型预热:在加载模型后,用一张小图(比如全黑图)先做一次推理,触发模型的初始化和CUDA内核编译,这样第一次正式推理就不会那么慢了。
- 使用硬件加速:确保PyTorch使用了CUDA,并且QT的渲染后端(如Windows上的ANGLE/Direct3D)也配置得当,以充分利用GPU。
6. 总结
走完这一趟,你应该已经掌握了用QT搭建SDPose-Wholebody本地可视化工具的基本脉络。从环境准备、界面设计,到核心的多线程推理框架,再到关键点的可视化绘制,每一步都是在解决一个具体的工程问题。
这个工具现在具备了基本的查看功能,但还有很多可以打磨的地方,比如更完善的交互编辑、批量处理图片、导出JSON格式的关键点数据、与ControlNet等生成模型联动等等。这些都可以作为你后续扩展的方向。
开发过程中,最关键的可能是对SDPose-OOD官方推理代码的封装,你需要仔细阅读其gradio_app或inference.py部分的代码,理解它如何加载模型、预处理图像、运行推理、后处理得到关键点。把这部分封装成一个干净的SDPoseInference类,是整个工具能跑起来的基础。
希望这篇文章能为你提供一个清晰的起点。动手试一试,当你看到自己写的程序成功地在复杂的图片上勾勒出精准的人体姿态时,那种成就感一定会很棒。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。