目录一、分割检测二、图像预处理二、推理三、后处理与可视化3.1、后处理3.2、mask可视化四、完整pytorch代码一、分割检测注本篇只是阐述推理流程tensorrt实现后续跟进。yolov8-pose的tensorrt部署代码稍后更新还是在仓库GitHub - FeiYull/TensorRT-Alpha: TensorRT-Alpha supports YOLOv8、YOLOv7、YOLOv6、YOLOv5、YOLOv4、v3、YOLOX、YOLOR...CUDA IS ALL YOU NEED.It also supports end2end CUDA C acceleration and multi-batch inference.也可以关注TensorRT系列教程-CSDN博客以下是官方预测代码from ultralytics import YOLO model YOLO(modelyolov8n-pose.pt) model.predict(sourced:/Data/1.jpg, saveTrue)推理过程无非是图像预处理 - 推理 - 后处理 可视化这三个关键步骤在文件大概247行D:\CodePython\ultralytics\ultralytics\engine\predictor.py代码如下# Preprocess with profilers[0]: im self.preprocess(im0s) # 图像预处理 # Inference with profilers[1]: preds self.inference(im, *args, **kwargs) # 推理 # Postprocess with profilers[2]: self.results self.postprocess(preds, im, im0s) # 后处理二、图像预处理通过debug进入上述self.preprocess函数看到代码实现如下。处理流程大概是padding满足矩形推理图像通道转换即BGR装RGB检查图像数据是否连续存储顺序有HWC转为CHW然后归一化。需要注意原始pytorch框架图像预处理的时候会将图像缩放padding为HxW的图像其中H、W为32倍数而导出tensorrt的时候为了高效推理H、W 固定为640x640。def preprocess(self, im): Prepares input image before inference. Args: im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list. not_tensor not isinstance(im, torch.Tensor) if not_tensor: im np.stack(self.pre_transform(im)) im im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im np.ascontiguousarray(im) # contiguous im torch.from_numpy(im) img im.to(self.device) img img.half() if self.model.fp16 else img.float() # uint8 to fp16/32 if not_tensor: img / 255 # 0 - 255 to 0.0 - 1.0 return img二、推理图像预处理之后直接推理就行了这里是基于pytorch推理。def inference(self, im, *args, **kwargs): visualize increment_path(self.save_dir / Path(self.batch[0][0]).stem, mkdirTrue) if self.args.visualize and (not self.source_type.tensor) else False return self.model(im, augmentself.args.augment, visualizevisualize)三、后处理与可视化3.1、后处理640x640输入之后有两个输出其中output1尺寸为116X8400其中11648032,32为seg部分特征经过NMS之后输出为N*38其中384 2 32output2尺寸为32x160x160拿上面NMS后的特征图后面即N*38矩阵后面部分N*32的特征图和output2作矩阵乘法得到N*160*160的矩阵接着执行sigmiod然后拉平得到N*160*160 的mask。然后将bbox缩放160*160的坐标系如下代码用于截断越界的mask就是如下函数。最后将所有mask上采样到640*640然后用阀值0.5过一下。最后mask中只有0和1了结束。有关def crop_mask(masks, boxes):的理解def crop_mask(masks, boxes): It takes a mask and a bounding box, and returns a mask that is cropped to the bounding box Args: masks (torch.Tensor): [n, h, w] tensor of masks boxes (torch.Tensor): [n, 4] tensor of bbox coordinates in relative point form Returns: (torch.Tensor): The masks are being cropped to the bounding box. n, h, w masks.shape x1, y1, x2, y2 torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(n,1,1) r torch.arange(w, devicemasks.device, dtypex1.dtype)[None, None, :] # rows shape(1,1,w) c torch.arange(h, devicemasks.device, dtypex1.dtype)[None, :, None] # cols shape(1,h,1) return masks * ((r x1) * (r x2) * (c y1) * (c y2))上面代码最后一句return如下图理解mask中所有点例如点rc必须在bbox内部。做法就是将bbox缩放到和mask一样的坐标系160x160如下图然后使用绿色的bbox将mask进行截断3.2、mask可视化直接将mask从灰度图转为彩色图然后将类别对应的颜色乘以0.4最后加在彩色图上就行了。​四、完整pytorch代码将以上流程合并起来并加以修改完整代码如下import torch import cv2 as cv import numpy as np from ultralytics.data.augment import LetterBox from ultralytics.utils import ops from ultralytics.engine.results import Results import copy # path d:/Data/1.jpg path d:/Data/640640.jpg device cuda:0 conf 0.25 iou 0.7 # preprocess im cv.imread(path) # letterbox im [im] orig_imgs copy.deepcopy(im) im [LetterBox([640, 640], autoTrue, stride32)(imagex) for x in im] im im[0][None] # im np.stack(im) im im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im np.ascontiguousarray(im) # contiguous im torch.from_numpy(im) img im.to(device) img img.float() img / 255 # load model pt ckpt torch.load(yolov8n-seg.pt, map_locationcpu) model ckpt[model].to(device).float() # FP32 model model.eval() # inference preds model(img) # poseprocess p ops.non_max_suppression(preds[0], conf, iou, agnosticFalse, max_det300, nc80, classesNone) results [] # 如果导出onnx第二个输出维度是1应该就是mask需要后续上采样 proto preds[1][-1] if len(preds[1]) 3 else preds[1] # second output is len 3 if pt, but only 1 if exported??????? for i, pred in enumerate(p): orig_img orig_imgs[i] if not len(pred): # save empty boxes results.append(Results(orig_imgorig_img, pathpath, namesmodel.names, boxespred[:, :6])) continue masks ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsampleTrue) # HWC if not isinstance(orig_imgs, torch.Tensor): pred[:, :4] ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) results.append(Results(orig_imgorig_img, pathpath, namesmodel.names, boxespred[:, :6], masksmasks)) # show plot_args {line_width: None,boxes: True,conf: True, labels: True} plot_args[im_gpu] img[0] result results[0] plotted_img result.plot(**plot_args) cv.imshow(plotted_img, plotted_img) cv.waitKey(0) cv.destroyAllWindows()​