赞
踩
一个不知名大学生,江湖人称菜狗
original author: jacky Li
Email : 3435673055@qq.com
Time of completion:2023.1.6
Last edited: 2023.1.6
前不久,ultralytics发布了一个yolov5 7.0版本,在这个版本中,隆重推出了yolov5的seg版本,也就是实例分割版本。
也是奔着吃瓜要热乎的原则,立马上手尝试了一下,发现这个版本的mask回归头设计,还真的很高效。事实上在mask AP上的指标也不错!关键是速度快,可以在CPU下跑到realtime,还要啥自行车呢?
时至2023年,终于迎来了在CPU上实时跑实例分割的时代。
实测了一下,在目标笔记多的时候,也可以以将近60ms的速度跑在CPU上,同时精度看起来也还不错。
单人场景下也可以跑的很快。
但是今天我们要做的事情不是就这样跑一下,而是要导出到ONNX,用ONNXRuntime来推理。
首先导出yolov5的实例分割模型:
python export.py --weights ~/Downloads/yolov5s-seg.pt --include onnx
这一步比较简单,导出来的onnx如下:
如果你嫌麻烦,也可以直接转到文末的代码专区,可以直接获取到onnx文件,并且推理脚本都有。
接下来就是拿到这个seg模型,用onnxruntime来推理。
你需要知道的是,在yolov5-seg模型中,使用协方差与原型图来进行mask的预测,也就是说,在你的pred里面,输出的维度实际上是117. 它包括:
4 + 1 + 80 + 32
其中32就是mask的协方差系数。
与此同时,还有另外一个分支输出掩码,通过协方差与掩码的乘积,就可以知道mask的具体输出。
接下里就是onnxruntime推理,直接上代码:
- def detect(image):
- im, ratio, (dw, dh) = letterbox(image, img_size, stride=32, auto=False)
- im /= 255.0
-
- print_shape(im)
-
- inputs = {model.get_inputs()[0].name: im}
- pred, proto = model.run(None, inputs)
- pred = pred.reshape([1, -1, 4 + 1 + classes + 32])
- pred = torch.as_tensor(pred)
- proto = torch.as_tensor(proto)
- print_shape(pred, proto)
-
- pred = non_max_suppression(
- pred, conf_thres, iou_thres, None, agnostic_nms, max_det=max_det, nm=32
- )
-
- # Process predictions
- for i, det in enumerate(pred): # per image
- if len(det):
- masks = process_mask(
- proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True
- ) # HWC
- det[:, :4] = scale_boxes(
- im.shape[2:], det[:, :4], image.shape
- ).round() # rescale boxes to im0 size
- segments = reversed(masks2segments(masks))
- segments = [
- scale_segments(im.shape[2:], x, image.shape, normalize=False)
- for x in segments
- ]
-
- boxes = det[:, :4]
- labels = det[:, 5]
- scores = det[:, 4]
- # print(boxes)
- # print(scores)
- image = visualize_det_cv2_part(image, scores, labels, boxes)
- image = vis_segments_mask(image, segments)
- return image

如果需要代码,请私聊博主,博主看见回。
如果感觉博主讲的对您有用,请点个关注支持一下吧,将会对此类问题持续更新……
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。