基于yolov8(ultralytics)工程导出的FastSAM的onnx模型,后处理和yolov8seg是一样的。
模型和完整测试代码。
1 FastSAM 导出 onnx
导出onnx的方式有两种,一种使用FastSAM工程,一种是使用yolov8(ultralytics)工程。本篇博客使用yolov8工程进行导出onnx。导出FastSAM和导出yolov8seg需要修改的地方一样的。本示例使用的是FastSAM-s,效果不是很好。需要修改两个地方。
第一处修改:
# 导出 onnx 增加
y = []
for i in range(self.nl):
t1 = self.cv2[i](x[i])
t2 = self.cv3[i](x[i])
y.append(t1)
y.append(t2)
return y
第二处修改:
# 导出 onnx 增加(修改)
# mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2) # mask coefficients
mc = [self.cv4[i](x[i]) for i in range(self.nl)]
x = self.detect(self, x)
return x, mc, p
增加保存onnx代码:
print("=========== onnx =========== ")
import torch
dummy_input = torch.randn(1, 3, 640, 640)
input_names = ["data"]
output_names = ["cls1", "reg1", "cls2", "reg2", "cls3", "reg3", "mc1", "mc2", "mc3", "seg"]
torch.onnx.export(self.model, dummy_input, "./yolov8nseg_relu_80class_dfl.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=12)
print("======================== convert onnx Finished! .... ")
修改完以上运行如下代码:
from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPrompt
model = FastSAM('./weights/FastSAM-s.pt')
image_path = './images/test.jpg'
everything_results = model(image_path, retina_masks=True, imgsz=640, conf=0.4, iou=0.9)
2 onnx 测试效果
3 RKNN板端测试效果
rknn仿真测试效果
板端实际效果(颜色配的不是很好,凑合看)
模型输入分辨率640x640,使用芯片rk3588。