最近在一家小公司做实习，领导想要使用mask rcnn在边缘设备上进行实例分割任务，让我训练pytorch版的模型，并将生成的.pt格式模型转成onnx格式，以备后用。
本次实践达到的效果：重新训练mask rcnn模型，得到mask_rcnn.pth模型；然后使用.py文件转格式，得到mask_rcnn.onnx模型。之后想要在香橙派上转格式，出现算子Loop问题。所以，本次只涉及mask rcnn的重新训练和转onnx格式模型。
1. 前期准备与说明
本次训练是在Google Colaboratory上进行训练，主要是不需要配置相关软件版本，效率较高(不太了解Colab的小伙伴可以查看相关说明或者在本地部署相关环境)。本次mask rcnn模型的训练基于一个预训练好的模型(mask rcnn源项目代码地址放在文末 参考1)，使用本地的训练集和测试集，生成一个COCO2025文件，其中包括1.annotations，2.train2025，3.val2025。将其放在mask_rcnn目录下。
1.annotations当中包含instances_train2025.json和instances_val2025.json。
(以下内容是本人尝试的过程中出现的问题，不具有普适性，可以忽略)
之前生成的训练集和测试集的标签文件不能拿来直接用，需要做以下调整：
1) 批量删除参数"file_name":"images/xxx.jpg"中的'images/'，不然训练的时候找不到训练图片路径；
2）将"categories"下的{"id":0,...,"id":1}，改成"1"和"2"。(根据自己训练的类别做具体调整)
"categories":[
{
"id":1, # "id":0
"name":"goods"
},
{
"id":2, # "id":1
"name":"googs1"
}
]
3）批量修改"category_id":0 和"category_id":1，分别改成"category_id":1和"category_id":2。
2.在train2025和val2025目录下，分别放置训练图片和测试图片。如果数量太多，最好在本地打包压缩，上传压缩包，然后解压到指定目录(Colab上传文件的速度太慢)
得到这样的目录：
2.开始训练
1.新建.ipynb文件，将下载的项目打包成压缩包，上传到colab中；
2.挂载谷歌硬盘；(防止colab连接中断，永久挂载项目)
# 挂载Google Drive
from google.colab import drive
drive.mount("/content/drive",force_remount=True)
3.在打开的.ipynb中使用命令，解压文件(等待压缩文件完成上传后再执行以下脚本)；
import zipfile
# 定义zip文件路径，文件名以mask_rcnn.zip为例
zip_path = "/content/mask_rcnn.zip"
# 定义解压目录文件夹
extract_path = "/content/drive/MyDrive"
# 解压文件
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(extract_path)
print("解压完成！")
4.切换到项目目录，并确认当前文件夹中的内容
# 切换目录
import os
os.chdir("/content/drive/MyDrive/mask_rcnn")
# 查看当前目录下的文件内容
!ls
5.执行训练
# 开始训练
# --num-classes=2 原来是90类，现在只有2类
!python train.py --num-classes=2
另外，这个项目还可以实现断点续传功能，使用参数 --resume 即可。示例：其中model_x中的x指的是具体的开始续传的epoch。
# 断点续传/接着训练参数，指定开始继续训练的模型地址以及当前的epoch
!python train.py --num-classes=2 --resume="/content/drive/MyDrive/mask_rcnn/save_weights/model_x.pth"
补充说明：需要自行下载权重文件(resnet50.pth 参考2 ；maskrcnn_resnet50_fpn_coco.pth 参考3)并改名，将其放在mask_rcnn目录下；
在train.py中修改文件存放地址：weights_dict =
torch.load("/content/drive/MyDrive/mask_rcnn/maskrcnn_resnet50_fpn_coco.pth", map_location="cpu")；
指定训练/测试数据存放地址：parser.add_argument('--data-path', default='/content/drive/MyDrive/mask_rcnn/COCO2025', help='dataset')；
3.执行预测
本人在执行的过程中，发现pillow版本太高了，需要下载较低版本的pillow。本人使用了7.1.2版本，实现模型测试没有问题。
# 下载指定版本的pillow
!pip install pillow=7.1.2
然后开始预测模型的效果。
# 预测
!python predict.py
另外说明：执行测试的过程中，需要指定测试图片的地址，本人将原图放在了"/content/drive/MyDrive/mask_rcnn /picture1"目录下，将结果放在了"/content/drive/MyDrive/mask_rcnn/test_result"目录下。
注意：需要修改predict.py中的参数，包括num_classes=2(更好地实现方式是，以命令行的形式指定参数)；
-
指定训练好的权重文件地址：weights_path="/content/drive/MyDrive/mask_rcnn/ save_weights/model_25.pth"；
-
测试图片地址：img_path="/content/drive/MyDrive/mask_rcnn/picture1/xxx.jpg"；
-
标签：label_json_path = "/content/drive/MyDrive/mask_rcnn/coco91_indices.json"；
-
修改保存测试图片结果的地址："plot_img.save("/content/drive/MyDrive/mask_rcnn/ test_result/xxx_测试结果.jpg")"。
plot_img.save("/content/drive/MyDrive/mask_rcnn/test/xxx.jpg")
4.模型转换
将步骤2 生成的model_x.pth转换成.onnx模型。下面是转格式的代码：
# 加载预训练的Mask R-CNN模型
import torch
import torchvision
from torchvision.models.detection import maskrcnn_resnet50_fpn
import onnx
# 加载自己训练好的模型
# 此处设置参数pretrained=False表示不使用官方训练好的模型，而使用自己的
# num_classes=3 表示2个目标分类，外加1个背景 (此处参数不修改一直报错)
model = maskrcnn_resnet50_fpn(pretrained=False, num_classes=3)
pretrained_weights = torch.load('/content/drive/MyDrive/mask_rcnn/save_weights/model_25.pth')
# 使用到模型中的pretrained_weights["model"]
model.load_state_dict(pretrained_weights["model"])
model.eval()
# 此处大小根据自己图片大小的具体情况进行修改，修改后面两个参数
# (batch_size, channels, w, h)
dummy_input = torch.randn(1, 3, 1920, 1088)
# 导出模型到 ONNX 格式
onnx_model_path = "/content/drive/MyDrive/mask_rcnn/save_weights/mask_rcnn_model.onnx"
# 调用torch.onnx.export() 将pth格式转成onnx格式
torch.onnx.export(
model,
dummy_input, # dummy_input.to(device)
onnx_model_path, # 输出 onnx格式文件的名称
export_params=True, # 存储训练参数
opset_version=13, # 建议使用11或更高版本13
do_constant_folding=True, # 优化常量
input_names=['input'], # 输入节点名称
output_names=['boxes', 'labels', 'scores', 'masks'], # Mask R-CNN的输出节点名称
# dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}},
# dynamic_axes=None
dynamic_axes={
'input': {0: 'batch_size', 2: 'height', 3: 'width'},
'boxes': {0: 'num_boxes'},
'labels': {0: 'num_boxes'},
'scores': {0: 'num_boxes'},
'masks': {0: 'num_boxes'}
}
)
# 加载 ONNX 模型并验证
onnx_model = onnx.load(onnx_model_path)
onnx.checker.check_model(onnx_model)
# 使用 ONNX Runtime 进行推理
onnx_session = ort.InferenceSession(onnx_model_path)
# (1, 3, 1920, 1088)
x = torch.randn(1, 3, 1920, 1088).numpy()
onnx_output = onnx_session.run(output_names, {input_names[0]:x})[0]
print("模型已成功导出为ONNX格式!")
5.可视化验证onnx模型的有效性
下载相关依赖，包括：onnx,onnxruntime,onnxscript等。
# 基础ONNX库 (必需)
! pip install onnx
# PyTorch的ONNX扩展 (建议安装)
! pip install onnxscript
# ONNX运行环境 (可选，用于验证导出的模型)
! pip install onnxruntime
以下代码是看网上的一些资料，然后编写，可以正常运行。如果有其他特殊需求，请自行更改。
# 实例分割
import numpy as np
import cv2
import random
from google.colab.patches import cv2_imshow
import onnxruntime as ort
import onnx
import torchvision.transforms as transforms
def visualize_instances(image, boxes, labels, scores, masks, score_threshold=0.5):
"""
可视化Mask R-CNN的实例分割结果
参数:
image: 原始图像 (RGB格式)
boxes: 边界框坐标，形状为 [N, 4]，格式为 [x1, y1, x2, y2]
labels: 类别标签，形状为 [N]
scores: 置信度分数，形状为 [N]
masks: 掩码，形状为 [N, H, W]，其中N是实例数量
score_threshold: 置信度阈值，用于筛选结果
"""
# 创建图像副本用于绘制
result = image.copy()
h, w = image.shape[:2] # 图像的高和宽
# 筛选高置信度的结果
high_score_indices = np.where(scores > score_threshold)[0]
if len(high_score_indices) == 0:
return result
# 为每个实例生成随机颜色
colors = [
(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
for _ in range(len(high_score_indices))
]
# 处理每个实例
for i, idx in enumerate(high_score_indices):
# 获取实例数据
box = boxes[idx]
label = labels[idx]
score = scores[idx]
mask = masks[idx]
"""
在原始图像上叠加一个带有透明度的彩色蒙版，用于突出显示图像中的特定区域。
1.创建彩色蒙版
创建了一个与原始图像result尺寸相同的黑色蒙版overlay(形状为H*W*3的零矩阵)
2.处理掩码(mask)
将输入的掩码形状从(1,H,W) 压缩为(H,W);
使用阈值0.9将掩码转换为二值图像(大于阈值的区域设为1，其他设为0)
3.应用颜色到蒙版
在蒙版上，将二值掩码中值为1的区域设置为指定颜色 overlay_color;
得到 在指定区域有颜色，其他区域为黑色的蒙版的结果
4.叠加图像
使用cv2.addWeighted函数，将原始图像与彩色蒙版按一定透明度叠加;
alpha参数设置成0.5，表示蒙版的透明度为50%；
最终结果是原始图像上叠加了半透明的彩色区域，突出显示掩码标记的目标区域
"""
# 指定masks的填充颜色
overlay_color = (0, 255, 0) # 绿色
# 1. 创建彩色蒙版 (H,W,3)
overlay = np.zeros_like(result) # 形状 (H,W,3)
squeezed_mask = mask.squeeze(0) # 从(1,H,W)变为(H,W)
threshold = 0.9 # 根据模型调整阈值
binary_mask = (squeezed_mask > threshold).astype(np.uint8)
# 2. 处理mask维度并应用颜色
# overlay[squeezed_mask == 1] = overlay_color
# overlay[squeezed_mask != 0] = overlay_color
overlay[binary_mask == 1] = overlay_color
# 3. 设置透明度并叠加
alpha = 0.5
result = cv2.addWeighted(result, 1, overlay, alpha, 0)
# 绘制边界框
x1, y1, x2, y2 = map(int, box)
print(x1,y1,x2,y2)
cv2.rectangle(result, (x1, y1), (x2, y2), colors[i], 2)
# 绘制类别标签和置信度
label_text = f"Class {label}: {score:.2f}"
cv2.putText(result, label_text, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[i], 2)
return result
# 使用示例
if __name__ == "__main__":
# 读取图像
image = cv2.imread("/content/drive/MyDrive/mask_rcnn/picture1/xxx.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转换为RGB格式
img_tensor = transforms.ToTensor()(image).unsqueeze(0) # 添加batch维度
# 初始化ONNX模型
session = ort.InferenceSession('/content/drive/MyDrive/mask_rcnn/save_weights/mask_rcnn_model.onnx', providers=['CPUExecutionProvider'])
# ONNX模型推理
onnx_output = session.run(
['boxes', 'labels', 'scores', 'masks'],
{'input': img_tensor.numpy()}
)
# 模型输出
boxes = onnx_output[0]
labels = onnx_output[1]
scores = onnx_output[2]
masks = onnx_output[3]
# 可视化结果
result = visualize_instances(image, boxes, labels, scores, masks)
# 保存或显示结果
result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
cv2.imwrite("segmentation_result.jpg", result_bgr)
cv2_imshow(result_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()
本人发布以上内容，只是希望遇到类似问题的小伙伴学习参考，希望大家少踩一些坑。如有雷同，那是我的问题，请联系我。
参考： 1.github.com/WZMIAOMIAO/… 2.github.com/WZMIAOMIAO/… 3.github.com/WZMIAOMIAO/…
