Yolo26 模型转换 onnx 再转换 om 模型到 Ascend310B4 运行经过前言原来以为说什么 yolo26 转换成 om 模型不能在昇腾上面运行实际呢就是把训练好的模型转换好就行了这里只要昇腾芯片对 onnx 模型版本的要求一般 opset 11 较为稳定以及昇腾芯片能够适配的 CANN 版本就好了其他就是一个运行匹配的问题此次为什么能够花费这么久呢完全是因为自己不理解 yolo26 模型的输出是什么意思蒙起眼睛来打靶完全打不中后来参考了一个博主的把需求弄清楚就好了。版本对应问题这里再次强调一下版本我的芯片是Ascend310B4对应的HDK为23.0.rc3芯片适配的CANN版本是8.1.RC1查询CANN版本兼容性-CANN社区版8.1.RC1-昇腾社区适配的opset版本是 11 查询通过查看昇腾CANN对ONNX的支持 链接https://www.hiascend.com/document/detail/zh/canncommercial/900/API/aolapi/operatorlist_00177.html 通过查看ONNX对应的ONNX opset链接 https://runtime.onnx.org.cn/docs/reference/compatibility.html代码对于代码来说主要是一个理解转化的模型需要进行 NMS 就好了理解模型的输出就能够很好的解决问题。比如模型的输入是[batch,300,6]那么就要明白这个数据代表的意思是什么加入 batch 是 1 那么如果输入一张图片的话就会有 300 个框的结果6 代表的是 [x1, y1, x2, y2, score, class_id]参数含义x1左上角 x 坐标y1左上角 y 坐标x2右下角 x 坐标y2右下角 y 坐标score 模型认为这个框里“确实有目标”的概率class_id 表示这个目标属于哪个类别总结成一句话就是在图像的某个位置我认为那里有一个目标它的置信度是多少它属于哪个类别结果很明确一张图中不太可能有 300 个目标那么我们就需要对结果做筛选选出符合要求的目标此时就需要做 NMS 非极大值抑制非极大值抑制是关键因为 yolo26 的模型导出成 onnx 并没有直接给你抑制代码如下fromais_bench.infer.interfaceimportInferSession# 导入Ascend推理接口importcv2# 导入OpenCV库用于图像处理importnumpyasnp# 导入NumPy库用于数值计算frompathlibimportPath# 导入Path类用于路径操作# # 类别定义# CLASS_NAMES{0:frogman,# 类别ID 0: 蛙人1:auv,# 类别ID 1: 自主水下航行器2:chain,# 类别ID 2: 链条3:boat,# 类别ID 3: 船只4:cage,# 类别ID 4: 笼子5:ball# 类别ID 5: 球}IMG_EXTS{.jpg,.jpeg,.png,.bmp,.tif,.tiff,.webp}# 支持的图像文件扩展名集合# # Letterbox# defletterbox(img,new_shape640,color(114,114,114)): 将图像通过letterbox方式缩放到指定大小保持宽高比并填充边缘 参数: img: 输入图像 new_shape: 目标尺寸宽高相同默认640 color: 填充颜色默认灰色(114,114,114) 返回: img: 处理后的图像 r: 缩放比例 (dw, dh): 填充的宽度和高度未取整前的 h,wimg.shape[:2]# 获取原始图像高度和宽度rmin(new_shape/h,new_shape/w)# 计算缩放比例选择较小的缩放因子以确保图像能完整放入new_w,new_hint(w*r),int(h*r)# 计算缩放后的新宽度和新高度resizedcv2.resize(img,(new_w,new_h))# 将图像缩放到新尺寸dwnew_shape-new_w# 计算需要填充的宽度差值dhnew_shape-new_h# 计算需要填充的高度差值dw/2# 左右均分填充宽度dh/2# 上下均分填充高度top,bottomint(round(dh-0.1)),int(round(dh0.1))# 计算上边和下边的填充像素数left,rightint(round(dw-0.1)),int(round(dw0.1))# 计算左边和右边的填充像素数imgcv2.copyMakeBorder(# 添加边框填充resized,# 输入图像top,bottom,left,right,# 四边填充大小cv2.BORDER_CONSTANT,# 边框类型为常数填充valuecolor# 填充颜色值)returnimg,r,(dw,dh)# 返回处理后的图像、缩放比例和填充值# # NMS# defnms(boxes,scores,iou_thres0.45): 非极大值抑制算法去除重叠度高的冗余检测框 参数: boxes: 边界框坐标列表 [[x1,y1,x2,y2], ...] scores: 对应的置信度分数 iou_thres: IoU阈值高于此阈值的框会被抑制 返回: 保留的边界框索引列表 iflen(boxes)0:# 如果没有边界框直接返回空列表return[]xywh[]# 存储转换为[x,y,width,height]格式的边界框forx1,y1,x2,y2inboxes:# 遍历所有边界框xywh.append([x1,y1,x2-x1,y2-y1])# 转换为[x,y,宽度,高度]格式idxscv2.dnn.NMSBoxes(# 调用OpenCV的NMS函数xywh,# 边界框列表scores.tolist(),# 置信度分数列表0.0,# 置信度阈值此处设为0因为之前已经过滤过iou_thres# IoU阈值)iflen(idxs)0:# 如果没有保留的框return[]returnidxs.flatten()# 将索引展平为一维数组返回# # 后处理YOLOv26: 1,300,6# defpostprocess(pred,orig_shape,ratio,pad,conf_thres0.25,iou_thres0.45): 模型输出后处理置信度过滤、NMS、坐标还原 参数: pred: 模型预测输出形状为(300,6)每行格式[x1,y1,x2,y2,score,class_id] orig_shape: 原始图像尺寸(高度,宽度) ratio: 缩放比例 pad: 填充值(dw,dh) conf_thres: 置信度阈值 iou_thres: NMS的IoU阈值 返回: boxes: 过滤并还原后的边界框 scores: 对应的置信度分数 cls_ids: 对应的类别ID orig_h,orig_worig_shape# 获取原始图像高度和宽度dw,dhpad# 获取填充值prednp.squeeze(pred)# 去除批次维度从(1,300,6)变为(300,6)boxespred[:,:4]# 提取边界框坐标 [x1,y1,x2,y2]scorespred[:,4]# 提取置信度分数cls_idspred[:,5].astype(int)# 提取类别ID并转换为整数类型# 置信度过滤保留置信度大于阈值的检测结果maskscoresconf_thres boxes,scores,cls_idsboxes[mask],scores[mask],cls_ids[mask]iflen(boxes)0:# 如果没有检测结果return[],[],[]# 非极大值抑制(NMS)keepnms(boxes,scores,iou_thres)boxesboxes[keep]# 保留NMS后的边界框scoresscores[keep]# 保留对应的置信度cls_idscls_ids[keep]# 保留对应的类别ID# 还原坐标将letterbox后的坐标映射回原始图像坐标boxes[:,[0,2]](boxes[:,[0,2]]-dw)/ratio# 还原x1和x2坐标boxes[:,[1,3]](boxes[:,[1,3]]-dh)/ratio# 还原y1和y2坐标# 坐标裁剪确保边界框不超出原始图像范围boxes[:,[0,2]]boxes[:,[0,2]].clip(0,orig_w)# 裁剪x坐标到[0,宽度]boxes[:,[1,3]]boxes[:,[1,3]].clip(0,orig_h)# 裁剪y坐标到[0,高度]returnboxes,scores,cls_ids# 返回处理后的结果# # 画框已修复OpenCV报错# defdraw_boxes(img,boxes,scores,cls_ids): 在图像上绘制检测框、标签和置信度 参数: img: 输入图像 boxes: 边界框坐标列表 scores: 置信度分数列表 cls_ids: 类别ID列表 返回: 绘制了检测结果的图像 forbox,score,cidinzip(boxes,scores,cls_ids):# 遍历每个检测结果x1,y1,x2,y2map(int,box)# 将坐标转换为整数类型labelf{CLASS_NAMES.get(int(cid),cid)}{score:.2f}# 生成标签文本类别名置信度# ⚠️ 关键修复必须强制 int避免 OpenCV 报错# 根据类别ID动态生成颜色确保值在0-255范围内且为整数color(int((37*int(cid))%255),# 红色分量int((17*int(cid))%255),# 绿色分量int((29*int(cid))%255)# 蓝色分量)# 绘制边界框矩形cv2.rectangle(img,(x1,y1),(x2,y2),color,2)# 线条粗细为2像素# 获取标签文本的尺寸(tw,th),_cv2.getTextSize(label,cv2.FONT_HERSHEY_SIMPLEX,0.6,2)# 计算标签背景矩形的上边界避免超出图像顶部y_topmax(0,y1-th-6)# 绘制标签背景矩形cv2.rectangle(img,(x1,y_top),(x1tw,y1),color,-1)# -1表示填充矩形# 绘制标签文本cv2.putText(img,# 目标图像label,# 文本内容(x1,max(0,y1-4)),# 文本位置cv2.FONT_HERSHEY_SIMPLEX,# 字体类型0.6,# 字体大小(255,255,255),# 文本颜色白色2,# 文本粗细cv2.LINE_AA# 抗锯齿线条类型)returnimg# 返回绘制后的图像# # 获取图片# defget_images(folder:Path): 递归获取文件夹下所有支持的图像文件 参数: folder: 文件夹路径Path对象 返回: 图像文件路径列表 imgs[]# 存储图像路径的列表forpinsorted(folder.rglob(*)):# 递归遍历文件夹下所有文件排序后ifp.suffix.lower()inIMG_EXTS:# 如果文件扩展名在支持的图像格式集合中imgs.append(p)# 添加到列表returnimgs# # 主函数# defmain():主函数解析命令行参数、加载模型、执行推理并保存结果importargparse# 导入命令行参数解析模块parserargparse.ArgumentParser()# 创建参数解析器parser.add_argument(--model,typestr,requiredTrue)# 模型文件路径必需parser.add_argument(--source,typestr,requiredTrue)# 输入图像源文件夹路径必需parser.add_argument(--save,typestr,default./output)# 输出保存目录默认为./outputparser.add_argument(--device,typeint,default0)# 昇腾设备ID默认为0parser.add_argument(--conf,typefloat,default0.20)# 置信度阈值默认0.20parser.add_argument(--iou,typefloat,default0.40)# NMS的IoU阈值默认0.40argsparser.parse_args()# 解析命令行参数# # load model# sessionInferSession(args.device,args.model)# 创建推理会话加载模型到指定设备save_dirPath(args.save)# 创建输出保存目录的Path对象img_dirsave_dir/images# 保存绘制后图像的子目录label_dirsave_dir/labels# 保存YOLO格式标签的子目录img_dir.mkdir(parentsTrue,exist_okTrue)# 创建图像保存目录如果父目录不存在则递归创建label_dir.mkdir(parentsTrue,exist_okTrue)# 创建标签保存目录# # images# imagesget_images(Path(args.source))# 获取所有待处理的图像文件print(f[INFO] Found{len(images)}images)# 打印找到的图像数量# # inference loop# forimg_pathinimages:# 遍历每个图像文件imgcv2.imread(str(img_path))# 读取图像BGR格式ifimgisNone:# 如果图像读取失败continue# 跳过该图像h,wimg.shape[:2]# 获取原始图像的高度和宽度# preprocess预处理img_lb,ratio,padletterbox(img,640)# 对图像进行letterbox处理缩放到640x640img_rgbcv2.cvtColor(img_lb,cv2.COLOR_BGR2RGB)# 将BGR格式转换为RGB格式inpimg_rgb.astype(np.float32)/255.0# 归一化到[0,1]范围并转换为float32类型inpnp.transpose(inp,(2,0,1))# 调整维度顺序从(H,W,C)转换为(C,H,W)inpnp.expand_dims(inp,0)# 添加批次维度变为(1,C,H,W)inpnp.ascontiguousarray(inp)# 确保数组在内存中是连续的# inference推理outsession.infer([inp])[0]# 执行推理获取输出假设第一个输出是检测结果predout[0]# 获取预测结果形状为(300,6)# postprocess后处理boxes,scores,cls_idspostprocess(pred,# 预测结果(h,w),# 原始图像尺寸ratio,# 缩放比例pad,# 填充值args.conf,# 置信度阈值args.iou# IoU阈值)# # save YOLO txt保存YOLO格式标签# txt_pathlabel_dir/f{img_path.stem}.txt# 生成标签文件路径与图像同名扩展名为.txtlines[]# 存储标签行的列表forbox,score,cidinzip(boxes,scores,cls_ids):# 遍历每个检测结果x1,y1,x2,y2box# 获取边界框坐标# 转换为YOLO格式中心点坐标和宽高均归一化到[0,1]cx(x1x2)/2/w# 中心点x坐标归一化cy(y1y2)/2/h# 中心点y坐标归一化bw(x2-x1)/w# 边界框宽度归一化bh(y2-y1)/h# 边界框高度归一化# 格式类别ID 中心x 中心y 宽度 高度 置信度lines.append(f{int(cid)}{cx:.6f}{cy:.6f}{bw:.6f}{bh:.6f}{score:.6f})txt_path.write_text(\n.join(lines))# 将标签写入文件每行一个检测结果# # save image保存绘制了检测框的图像# visdraw_boxes(img.copy(),boxes,scores,cls_ids)# 在图像副本上绘制检测结果save_img_pathimg_dir/f{img_path.stem}.jpg# 生成保存图像路径转换为jpg格式cv2.imwrite(str(save_img_path),vis)# 保存绘制后的图像print(f{img_path.name}:{len(boxes)}objects)# 打印检测到的目标数量if__name____main__:# 如果作为主程序运行main()# 调用主函数
Yolo26 模型转换 onnx 再转换 om 模型到 Ascend310B4 运行经过
发布时间：2026/6/9 7:44:09
Yolo26 模型转换 onnx 再转换 om 模型到 Ascend310B4 运行经过前言原来以为说什么 yolo26 转换成 om 模型不能在昇腾上面运行实际呢就是把训练好的模型转换好就行了这里只要昇腾芯片对 onnx 模型版本的要求一般 opset 11 较为稳定以及昇腾芯片能够适配的 CANN 版本就好了其他就是一个运行匹配的问题此次为什么能够花费这么久呢完全是因为自己不理解 yolo26 模型的输出是什么意思蒙起眼睛来打靶完全打不中后来参考了一个博主的把需求弄清楚就好了。版本对应问题这里再次强调一下版本我的芯片是Ascend310B4对应的HDK为23.0.rc3芯片适配的CANN版本是8.1.RC1查询CANN版本兼容性-CANN社区版8.1.RC1-昇腾社区适配的opset版本是 11 查询通过查看昇腾CANN对ONNX的支持 链接https://www.hiascend.com/document/detail/zh/canncommercial/900/API/aolapi/operatorlist_00177.html 通过查看ONNX对应的ONNX opset链接 https://runtime.onnx.org.cn/docs/reference/compatibility.html代码对于代码来说主要是一个理解转化的模型需要进行 NMS 就好了理解模型的输出就能够很好的解决问题。比如模型的输入是[batch,300,6]那么就要明白这个数据代表的意思是什么加入 batch 是 1 那么如果输入一张图片的话就会有 300 个框的结果6 代表的是 [x1, y1, x2, y2, score, class_id]参数含义x1左上角 x 坐标y1左上角 y 坐标x2右下角 x 坐标y2右下角 y 坐标score 模型认为这个框里“确实有目标”的概率class_id 表示这个目标属于哪个类别总结成一句话就是在图像的某个位置我认为那里有一个目标它的置信度是多少它属于哪个类别结果很明确一张图中不太可能有 300 个目标那么我们就需要对结果做筛选选出符合要求的目标此时就需要做 NMS 非极大值抑制非极大值抑制是关键因为 yolo26 的模型导出成 onnx 并没有直接给你抑制代码如下fromais_bench.infer.interfaceimportInferSession# 导入Ascend推理接口importcv2# 导入OpenCV库用于图像处理importnumpyasnp# 导入NumPy库用于数值计算frompathlibimportPath# 导入Path类用于路径操作# # 类别定义# CLASS_NAMES{0:frogman,# 类别ID 0: 蛙人1:auv,# 类别ID 1: 自主水下航行器2:chain,# 类别ID 2: 链条3:boat,# 类别ID 3: 船只4:cage,# 类别ID 4: 笼子5:ball# 类别ID 5: 球}IMG_EXTS{.jpg,.jpeg,.png,.bmp,.tif,.tiff,.webp}# 支持的图像文件扩展名集合# # Letterbox# defletterbox(img,new_shape640,color(114,114,114)): 将图像通过letterbox方式缩放到指定大小保持宽高比并填充边缘 参数: img: 输入图像 new_shape: 目标尺寸宽高相同默认640 color: 填充颜色默认灰色(114,114,114) 返回: img: 处理后的图像 r: 缩放比例 (dw, dh): 填充的宽度和高度未取整前的 h,wimg.shape[:2]# 获取原始图像高度和宽度rmin(new_shape/h,new_shape/w)# 计算缩放比例选择较小的缩放因子以确保图像能完整放入new_w,new_hint(w*r),int(h*r)# 计算缩放后的新宽度和新高度resizedcv2.resize(img,(new_w,new_h))# 将图像缩放到新尺寸dwnew_shape-new_w# 计算需要填充的宽度差值dhnew_shape-new_h# 计算需要填充的高度差值dw/2# 左右均分填充宽度dh/2# 上下均分填充高度top,bottomint(round(dh-0.1)),int(round(dh0.1))# 计算上边和下边的填充像素数left,rightint(round(dw-0.1)),int(round(dw0.1))# 计算左边和右边的填充像素数imgcv2.copyMakeBorder(# 添加边框填充resized,# 输入图像top,bottom,left,right,# 四边填充大小cv2.BORDER_CONSTANT,# 边框类型为常数填充valuecolor# 填充颜色值)returnimg,r,(dw,dh)# 返回处理后的图像、缩放比例和填充值# # NMS# defnms(boxes,scores,iou_thres0.45): 非极大值抑制算法去除重叠度高的冗余检测框 参数: boxes: 边界框坐标列表 [[x1,y1,x2,y2], ...] scores: 对应的置信度分数 iou_thres: IoU阈值高于此阈值的框会被抑制 返回: 保留的边界框索引列表 iflen(boxes)0:# 如果没有边界框直接返回空列表return[]xywh[]# 存储转换为[x,y,width,height]格式的边界框forx1,y1,x2,y2inboxes:# 遍历所有边界框xywh.append([x1,y1,x2-x1,y2-y1])# 转换为[x,y,宽度,高度]格式idxscv2.dnn.NMSBoxes(# 调用OpenCV的NMS函数xywh,# 边界框列表scores.tolist(),# 置信度分数列表0.0,# 置信度阈值此处设为0因为之前已经过滤过iou_thres# IoU阈值)iflen(idxs)0:# 如果没有保留的框return[]returnidxs.flatten()# 将索引展平为一维数组返回# # 后处理YOLOv26: 1,300,6# defpostprocess(pred,orig_shape,ratio,pad,conf_thres0.25,iou_thres0.45): 模型输出后处理置信度过滤、NMS、坐标还原 参数: pred: 模型预测输出形状为(300,6)每行格式[x1,y1,x2,y2,score,class_id] orig_shape: 原始图像尺寸(高度,宽度) ratio: 缩放比例 pad: 填充值(dw,dh) conf_thres: 置信度阈值 iou_thres: NMS的IoU阈值 返回: boxes: 过滤并还原后的边界框 scores: 对应的置信度分数 cls_ids: 对应的类别ID orig_h,orig_worig_shape# 获取原始图像高度和宽度dw,dhpad# 获取填充值prednp.squeeze(pred)# 去除批次维度从(1,300,6)变为(300,6)boxespred[:,:4]# 提取边界框坐标 [x1,y1,x2,y2]scorespred[:,4]# 提取置信度分数cls_idspred[:,5].astype(int)# 提取类别ID并转换为整数类型# 置信度过滤保留置信度大于阈值的检测结果maskscoresconf_thres boxes,scores,cls_idsboxes[mask],scores[mask],cls_ids[mask]iflen(boxes)0:# 如果没有检测结果return[],[],[]# 非极大值抑制(NMS)keepnms(boxes,scores,iou_thres)boxesboxes[keep]# 保留NMS后的边界框scoresscores[keep]# 保留对应的置信度cls_idscls_ids[keep]# 保留对应的类别ID# 还原坐标将letterbox后的坐标映射回原始图像坐标boxes[:,[0,2]](boxes[:,[0,2]]-dw)/ratio# 还原x1和x2坐标boxes[:,[1,3]](boxes[:,[1,3]]-dh)/ratio# 还原y1和y2坐标# 坐标裁剪确保边界框不超出原始图像范围boxes[:,[0,2]]boxes[:,[0,2]].clip(0,orig_w)# 裁剪x坐标到[0,宽度]boxes[:,[1,3]]boxes[:,[1,3]].clip(0,orig_h)# 裁剪y坐标到[0,高度]returnboxes,scores,cls_ids# 返回处理后的结果# # 画框已修复OpenCV报错# defdraw_boxes(img,boxes,scores,cls_ids): 在图像上绘制检测框、标签和置信度 参数: img: 输入图像 boxes: 边界框坐标列表 scores: 置信度分数列表 cls_ids: 类别ID列表 返回: 绘制了检测结果的图像 forbox,score,cidinzip(boxes,scores,cls_ids):# 遍历每个检测结果x1,y1,x2,y2map(int,box)# 将坐标转换为整数类型labelf{CLASS_NAMES.get(int(cid),cid)}{score:.2f}# 生成标签文本类别名置信度# ⚠️ 关键修复必须强制 int避免 OpenCV 报错# 根据类别ID动态生成颜色确保值在0-255范围内且为整数color(int((37*int(cid))%255),# 红色分量int((17*int(cid))%255),# 绿色分量int((29*int(cid))%255)# 蓝色分量)# 绘制边界框矩形cv2.rectangle(img,(x1,y1),(x2,y2),color,2)# 线条粗细为2像素# 获取标签文本的尺寸(tw,th),_cv2.getTextSize(label,cv2.FONT_HERSHEY_SIMPLEX,0.6,2)# 计算标签背景矩形的上边界避免超出图像顶部y_topmax(0,y1-th-6)# 绘制标签背景矩形cv2.rectangle(img,(x1,y_top),(x1tw,y1),color,-1)# -1表示填充矩形# 绘制标签文本cv2.putText(img,# 目标图像label,# 文本内容(x1,max(0,y1-4)),# 文本位置cv2.FONT_HERSHEY_SIMPLEX,# 字体类型0.6,# 字体大小(255,255,255),# 文本颜色白色2,# 文本粗细cv2.LINE_AA# 抗锯齿线条类型)returnimg# 返回绘制后的图像# # 获取图片# defget_images(folder:Path): 递归获取文件夹下所有支持的图像文件 参数: folder: 文件夹路径Path对象 返回: 图像文件路径列表 imgs[]# 存储图像路径的列表forpinsorted(folder.rglob(*)):# 递归遍历文件夹下所有文件排序后ifp.suffix.lower()inIMG_EXTS:# 如果文件扩展名在支持的图像格式集合中imgs.append(p)# 添加到列表returnimgs# # 主函数# defmain():主函数解析命令行参数、加载模型、执行推理并保存结果importargparse# 导入命令行参数解析模块parserargparse.ArgumentParser()# 创建参数解析器parser.add_argument(--model,typestr,requiredTrue)# 模型文件路径必需parser.add_argument(--source,typestr,requiredTrue)# 输入图像源文件夹路径必需parser.add_argument(--save,typestr,default./output)# 输出保存目录默认为./outputparser.add_argument(--device,typeint,default0)# 昇腾设备ID默认为0parser.add_argument(--conf,typefloat,default0.20)# 置信度阈值默认0.20parser.add_argument(--iou,typefloat,default0.40)# NMS的IoU阈值默认0.40argsparser.parse_args()# 解析命令行参数# # load model# sessionInferSession(args.device,args.model)# 创建推理会话加载模型到指定设备save_dirPath(args.save)# 创建输出保存目录的Path对象img_dirsave_dir/images# 保存绘制后图像的子目录label_dirsave_dir/labels# 保存YOLO格式标签的子目录img_dir.mkdir(parentsTrue,exist_okTrue)# 创建图像保存目录如果父目录不存在则递归创建label_dir.mkdir(parentsTrue,exist_okTrue)# 创建标签保存目录# # images# imagesget_images(Path(args.source))# 获取所有待处理的图像文件print(f[INFO] Found{len(images)}images)# 打印找到的图像数量# # inference loop# forimg_pathinimages:# 遍历每个图像文件imgcv2.imread(str(img_path))# 读取图像BGR格式ifimgisNone:# 如果图像读取失败continue# 跳过该图像h,wimg.shape[:2]# 获取原始图像的高度和宽度# preprocess预处理img_lb,ratio,padletterbox(img,640)# 对图像进行letterbox处理缩放到640x640img_rgbcv2.cvtColor(img_lb,cv2.COLOR_BGR2RGB)# 将BGR格式转换为RGB格式inpimg_rgb.astype(np.float32)/255.0# 归一化到[0,1]范围并转换为float32类型inpnp.transpose(inp,(2,0,1))# 调整维度顺序从(H,W,C)转换为(C,H,W)inpnp.expand_dims(inp,0)# 添加批次维度变为(1,C,H,W)inpnp.ascontiguousarray(inp)# 确保数组在内存中是连续的# inference推理outsession.infer([inp])[0]# 执行推理获取输出假设第一个输出是检测结果predout[0]# 获取预测结果形状为(300,6)# postprocess后处理boxes,scores,cls_idspostprocess(pred,# 预测结果(h,w),# 原始图像尺寸ratio,# 缩放比例pad,# 填充值args.conf,# 置信度阈值args.iou# IoU阈值)# # save YOLO txt保存YOLO格式标签# txt_pathlabel_dir/f{img_path.stem}.txt# 生成标签文件路径与图像同名扩展名为.txtlines[]# 存储标签行的列表forbox,score,cidinzip(boxes,scores,cls_ids):# 遍历每个检测结果x1,y1,x2,y2box# 获取边界框坐标# 转换为YOLO格式中心点坐标和宽高均归一化到[0,1]cx(x1x2)/2/w# 中心点x坐标归一化cy(y1y2)/2/h# 中心点y坐标归一化bw(x2-x1)/w# 边界框宽度归一化bh(y2-y1)/h# 边界框高度归一化# 格式类别ID 中心x 中心y 宽度 高度 置信度lines.append(f{int(cid)}{cx:.6f}{cy:.6f}{bw:.6f}{bh:.6f}{score:.6f})txt_path.write_text(\n.join(lines))# 将标签写入文件每行一个检测结果# # save image保存绘制了检测框的图像# visdraw_boxes(img.copy(),boxes,scores,cls_ids)# 在图像副本上绘制检测结果save_img_pathimg_dir/f{img_path.stem}.jpg# 生成保存图像路径转换为jpg格式cv2.imwrite(str(save_img_path),vis)# 保存绘制后的图像print(f{img_path.name}:{len(boxes)}objects)# 打印检测到的目标数量if__name____main__:# 如果作为主程序运行main()# 调用主函数
