FastSAM 部署 rknn

基于yolov8(ultralytics)工程导出的FastSAM的onnx模型,后处理和yolov8seg是一样的。

模型和完整测试代码

1 FastSAM 导出 onnx

导出onnx的方式有两种,一种使用FastSAM工程,一种是使用yolov8(ultralytics)工程。本篇博客使用yolov8工程进行导出onnx。导出FastSAM和导出yolov8seg需要修改的地方一样的。本示例使用的是FastSAM-s,效果不是很好。需要修改两个地方。

第一处修改:

python 复制代码
        # 导出 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

第二处修改:

python 复制代码
        # 导出 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代码:

python 复制代码
        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! .... ")

修改完以上运行如下代码:

python 复制代码
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。

相关推荐
微风粼粼25 分钟前
程序员在线接单
java·jvm·后端·python·eclipse·tomcat·dubbo
云天徽上40 分钟前
【PaddleOCR】OCR表格识别数据集介绍,包含PubTabNet、好未来表格识别、WTW中文场景表格等数据,持续更新中......
python·ocr·文字识别·表格识别·paddleocr·pp-ocrv5
你怎么知道我是队长1 小时前
python-input内置函数
开发语言·python
叹一曲当时只道是寻常1 小时前
Python实现优雅的目录结构打印工具
python
hbwhmama2 小时前
python高级变量XIII
python
费弗里3 小时前
Python全栈应用开发利器Dash 3.x新版本介绍(3)
python·dash
dme.3 小时前
Javascript之DOM操作
开发语言·javascript·爬虫·python·ecmascript
加油吧zkf3 小时前
AI大模型如何重塑软件开发流程?——结合目标检测的深度实践与代码示例
开发语言·图像处理·人工智能·python·yolo
t_hj3 小时前
python规划
python
SDUERPANG3 小时前
三维目标检测|Iou3D 代码解读一
人工智能·目标检测·3d