登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 业界 ›基于onnxruntime结合PyQt快速搭建视觉原型Demo ...

基于onnxruntime结合PyQt快速搭建视觉原型Demo

嘀荼酴 7 天前
  我在日常工作中经常使用PyQt和onnxruntime来快速生产demo软件，用于展示和测试，这里，我将以Yolov12为例，展示一下我的方案。
  首先我们需要使用Yolov12训练一个模型，并export出Onnx文件，这个部分网络上有很多内容，可以使用ultralytics框架做这个事情，我在这里就不在赘述了，接下来的步骤直接从onnxruntime开始。
  此处你需要针对你的模型去编写一个基于Onnxruntime的推理类，包括前处理，后处理，可视化等部分，做到输入是图片，输出是结果，要保证算法代码和软件代码的独立性。这里是我参考ultralytics框架写的一个yolov12的推理类，用于实例分割的推理。
  1. import cv2
  2. import numpy as np
  3. import onnxruntime as ort
  4. import torch
  5. import yoloSeg.utils.ops as ops
  6. from yoloSeg.utils.results import Results
  9. class YOLOv12Seg:
  10.     """
  11.     YOLOv12 segmentation model for performing instance segmentation using ONNX Runtime.
  13.     This class implements a YOLOv12 instance segmentation model using ONNX Runtime for inference. It handles
  14.     preprocessing of input images, running inference with the ONNX model, and postprocessing the results to
  15.     generate bounding boxes and segmentation masks.
  16.     """
  18.     def __init__(self, onnx_model, classes, conf=0.25, iou=0.7, imgsz=640):
  19.         """
  20.         Initialize the instance segmentation model using an ONNX model.
  21.         """
  22.         self.session = ort.InferenceSession(
  23.             onnx_model,
  24.             providers=["CPUExecutionProvider"]
  25.             # if torch.cuda.is_available()
  26.             # else ["CPUExecutionProvider"],
  27.         )
  28.         self.imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz
  29.         self.classes = classes
  30.         self.conf = conf
  31.         self.iou = iou
  33.     def __call__(self, img):
  34.         """
  35.         Run inference on the input image using the ONNX model.
  36.         """
  37.         prep_img = self.preprocess(img, self.imgsz)
  38.         outs = self.session.run(None, {self.session.get_inputs()[0].name: prep_img})
  39.         return self.postprocess(img, prep_img, outs)
  41.     def letterbox(self, img, new_shape=(640, 640)):
  42.         """
  43.         Resize and pad image while maintaining aspect ratio.
  44.         """
  45.         shape = img.shape[:2]  # current shape [height, width]
  46.         # Scale ratio (new / old)
  47.         r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  48.         # Compute padding
  49.         new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  50.         dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # wh padding
  51.         if shape[::-1] != new_unpad:  # resize
  52.             img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  53.         top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  54.         left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  55.         img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
  56.         return img
  58.     def preprocess(self, img, new_shape):
  59.         """
  60.         Preprocess the input image before feeding it into the model.
  61.         """
  62.         img = self.letterbox(img, new_shape)
  63.         img = img[..., ::-1].transpose([2, 0, 1])[None]  # BGR to RGB, BHWC to BCHW
  64.         img = np.ascontiguousarray(img)
  65.         img = img.astype(np.float32) / 255  # Normalize to [0, 1]
  66.         return img
  68.     def postprocess(self, img, prep_img, outs):
  69.         """
  70.         Post-process model predictions to extract meaningful results.
  71.         """
  72.         preds, protos = [torch.from_numpy(p) for p in outs]
  73.         preds = ops.non_max_suppression(preds, self.conf, self.iou, nc=len(self.classes))
  74.         results = []
  75.         for i, pred in enumerate(preds):
  76.             pred[:, :4] = ops.scale_boxes(prep_img.shape[2:], pred[:, :4], img.shape)
  77.             masks = self.process_mask(protos[i], pred[:, 6:], pred[:, :4], img.shape[:2])
  78.             results.append(Results(img, path="", names=self.classes, boxes=pred[:, :6], masks=masks))
  79.         return results
  81.     def process_mask(self, protos, masks_in, bboxes, shape):
  82.         """
  83.         Process prototype masks with predicted mask coefficients to generate instance segmentation masks.
  84.         """
  85.         c, mh, mw = protos.shape  # CHW
  86.         masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw)  # Matrix multiplication
  87.         masks = ops.scale_masks(masks[None], shape)[0]  # Scale masks to original image size
  88.         masks = ops.crop_mask(masks, bboxes)  # Crop masks to bounding boxes
  89.         return masks.gt_(0.0)  # Convert to binary masks
  91.     def visualize_segmentation(self, image, results, alpha=0.95):
  92.         # 创建图像副本
  93.         visualization = image.copy()
  94.         # 获取预测结果
  95.         if isinstance(results, list):
  96.             result = results[0]  # 只取第一个结果
  97.         else:
  98.             result = results
  99.         # 检查是否有分割掩码
  100.         if hasattr(result, 'masks') and result.masks is not None:
  101.             # 获取边界框、类别和置信度
  102.             boxes = result.boxes.cpu().numpy()
  103.             masks = result.masks.data.cpu().numpy()
  104.             # 生成随机颜色
  105.             num_instances = len(boxes)
  106.             colors = np.random.randint(0, 255, size=(num_instances, 3), dtype=np.uint8)
  107.             # 遍历每个实例
  108.             for i in range(num_instances):
  109.                 confidence = boxes.conf[i]
  110.                 class_id = int(boxes.cls[i])
  111.                 # 获取掩码并调整大小以匹配原始图像
  112.                 mask = masks[i]
  113.                 if mask.shape[:2] != image.shape[:2]:
  114.                     mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
  115.                 # 创建着色掩码
  116.                 color_mask = np.zeros_like(image)
  117.                 mask_bool = (mask > 0).astype(bool)
  118.                 color = colors[i].tolist()
  119.                 color_mask[mask_bool] = color
  120.                 # 将掩码与原始图像混合
  121.                 visualization = cv2.addWeighted(visualization, 1.0, color_mask, alpha, 0)
  122.                 # # 绘制边界框
  123.                 # x1, y1, x2, y2 = map(int, boxes.xyxy[i])
  124.                 # cv2.rectangle(visualization, (x1, y1), (x2, y2), color, 2)
  125.                 # # 获取类别名称
  126.                 # class_name = self.classes[class_id]
  127.                 # # 显示类别名称和置信度
  128.                 # label = f"{class_name}: {confidence:.2f}"
  129.                 # text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
  130.                 # cv2.rectangle(visualization, (x1, y1 - text_size[1] - 5), (x1 + text_size[0], y1), color, -1)
  131.                 # cv2.putText(visualization, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
  132.         return visualization
复制代码
  完成算法模块的独立后，接下来就是基于PyQt的软件模块开发，此处之所以强调两者互相独立，也是为了提高软件代码后续的重复使用率，以后即使换了一个算法，只需要修改几行代码就可以快速的实现第二个demo软件的开发。那么因为我们是视觉项目，所以整个软件需要强调的是原图，结果图，以及一些相关的功能按钮，接下来我将提供一个软件模板和对应的注释，方便大家使用。
  1. import sys
  2. import os
  3. from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
  4.                             QPushButton, QLabel, QFileDialog, QTabWidget, QScrollArea,
  5.                             QSplitter, QMessageBox)
  6. from PyQt5.QtGui import QPixmap, QImage
  7. from PyQt5.QtCore import Qt, pyqtSignal, QThread
  8. import cv2
  9. import numpy as np
  10. from yoloSeg.model import YOLOv12Seg as SegAI
  13. class ImageProcessor(QThread):
  14.     """处理图像的线程，避免UI阻塞"""
  15.     detection_completed = pyqtSignal(np.ndarray)
  16.    
  17.     def __init__(self):
  18.         super().__init__()
  19.         self.image = None
  20.         self.ai = SegAI("weights/best.onnx", ["goldLine"], conf=0.2, iou=0.8)
  21.         
  22.     def set_images(self, image):
  23.         self.image = image
  24.         
  25.     def run(self):
  26.         result_image = self.detection_algorithm()
  27.         self.detection_completed.emit(result_image)
  28.    
  29.     def detection_algorithm(self):
  30.         results = self.ai(self.image)
  31.         visualization = self.ai.visualize_segmentation(self.image, results)
  32.         return visualization
  35. class ImageViewer(QLabel):
  36.     """图像查看器组件"""
  37.     def __init__(self):
  38.         super().__init__()
  39.         self.setAlignment(Qt.AlignCenter)
  40.         self.setMinimumSize(300, 300)
  41.         self.setStyleSheet("border: 1px solid gray; background-color: #f0f0f0;")
  42.         self.setScaledContents(False)
  43.         self.current_pixmap = None
  44.         
  45.     def setImage(self, image):
  46.         if isinstance(image, np.ndarray):
  47.             # 将OpenCV图像转换为QPixmap
  48.             height, width, channel = image.shape
  49.             bytesPerLine = 3 * width
  50.             qImg = QImage(image.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
  51.             pixmap = QPixmap.fromImage(qImg)
  52.         else:
  53.             pixmap = QPixmap(image)
  54.             
  55.         if pixmap.isNull():
  56.             return
  57.         self.current_pixmap = pixmap
  58.         self.updatePixmap()
  59.             
  60.     def updatePixmap(self):
  61.         if self.current_pixmap:
  62.             # 保持纵横比缩放到标签大小
  63.             scaled_pixmap = self.current_pixmap.scaled(
  64.                 self.width(), self.height(),
  65.                 Qt.KeepAspectRatio, Qt.SmoothTransformation
  66.             )
  67.             super().setPixmap(scaled_pixmap)
  68.    
  69.     def resizeEvent(self, event):
  70.         self.updatePixmap()
  71.         super().resizeEvent(event)
  74. class ImageProcessingApp(QMainWindow):
  75.     def __init__(self):
  76.         super().__init__()
  77.         self.image = None
  78.         # 组件
  79.         self.initUI()
  80.         self.image_processor = ImageProcessor()
  81.         self.image_processor.detection_completed.connect(self.on_detection_completed)
  82.         
  83.     def initUI(self):
  84.         # 设置窗口标题和大小
  85.         self.setWindowTitle('检测软件')
  86.         self.setGeometry(100, 100, 1200, 700)
  87.         # 创建中央部件和总体布局
  88.         central_widget = QWidget()
  89.         self.setCentralWidget(central_widget)
  90.         main_layout = QVBoxLayout(central_widget)
  91.         # 创建顶部按钮布局
  92.         button_layout = QHBoxLayout()
  93.         # 添加按钮
  94.         self.btn_open = QPushButton('打开图片文件')
  95.         self.btn_detection = QPushButton('开始检测')
  96.         self.btn_save= QPushButton('保存图片')
  97.         # 添加按钮到布局
  98.         button_layout.addWidget(self.btn_open)
  99.         button_layout.addWidget(self.btn_detection)
  100.         button_layout.addWidget(self.btn_save)
  101.         # 连接按钮信号
  102.         self.btn_open.clicked.connect(self.open_file)
  103.         self.btn_detection.clicked.connect(self.detect_image)
  104.         self.btn_save.clicked.connect(self.save_image)
  105.         # 禁用未加载图片前的按钮
  106.         self.btn_save.setEnabled(False)
  107.         # 添加分隔器分割左右区域
  108.         splitter = QSplitter(Qt.Horizontal)
  109.         # 左侧：图片展示区
  110.         self.image_tabs = QTabWidget()
  111.         self.image_viewer = ImageViewer()
  112.         self.image_tabs.addTab(self.image_viewer, "图片队列")
  113.         splitter.addWidget(self.image_tabs)
  114.         # 右侧：结果显示区（使用标签页）
  115.         self.result_tabs = QTabWidget()
  116.         # 添加标签页
  117.         self.detection_tab = ImageViewer()
  118.         self.result_tabs.addTab(self.detection_tab, "图片检测")
  119.         splitter.addWidget(self.result_tabs)
  120.         # 设置分隔器比例
  121.         splitter.setSizes([600, 600])
  122.         # 添加按钮区域和分隔器到主布局
  123.         main_layout.addLayout(button_layout)
  124.         main_layout.addWidget(splitter, 1)
  125.         # 添加状态栏
  126.         self.statusBar().showMessage('就绪')
  127.         
  128.     def open_file(self):
  129.         """打开文件夹并加载图像"""
  130.         file_path, filetype = QFileDialog.getOpenFileName(self,  
  131.                                     "选取文件",  
  132.                                     os.getcwd(), # 起始路径
  133.                                     "Image Files (*.jpg *.jpeg *.png *.bmp *.tif *.tiff)")
  134.         if file_path:
  135.             self.image = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), -1)
  136.             # 显示图片
  137.             if self.image.any():
  138.                 self.image_viewer.setImage(self.image)
  139.             else:
  140.                 QMessageBox.warning(self, '警告', '无法加载任何图像文件！')
  141.    
  142.     def detect_image(self):
  143.         """检测图片"""
  144.         self.statusBar().showMessage('正在进行图片检测...')
  145.         # 使用线程处理图片检测
  146.         self.image_processor.set_images(self.image)
  147.         self.image_processor.start()
  148.    
  149.     def on_detection_completed(self, result_image):
  150.         """检测完成后的回调"""
  151.         self.detection_tab.setImage(result_image)
  152.         self.result_tabs.setCurrentIndex(1)  # 切换到图片检测标签页
  153.         # 保存检测结果
  154.         self.detection_result = result_image
  155.         self.statusBar().showMessage('图片检测完成')
  156.         self.btn_save.setEnabled(True) # 启用保存按钮
  158.     def save_image(self):
  159.         """保存最新阶段的图片"""
  160.         cv2.imwrite("result.jpg", self.detection_result)
  161.         self.statusBar().showMessage('图片保存完成')
  164. if __name__ == "__main__":
  165.     app = QApplication(sys.argv)
  166.     ex = ImageProcessingApp()
  167.     ex.show()
  168.     sys.exit(app.exec_())
  169.    
复制代码
   这里注意ImageProcessor类其实就是算法类在软件中的代理类，类似于一个协议，后续只需要按需求修改这个类和算法类即可。

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

使用道具 举报
相关推荐
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
嘀荼酴
7 天前

0

粉丝关注

10

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫