.pth转.weights/openCV-python + YOLO v3实现目标检测


前言

  毕设做的是水面目标的目标检测,因为要用Tkinter制作用户界面,于是采用openCV库来实现图像的处理,恰好openCV支持YOLO v3的部署……


一、基于Pytorch框架的YOLO v3

  因为对Python比较熟悉,综合考虑后选择用Pytorch框架实现,奈何自己太菜,只能去Github上找现成的YOLO v3代码,这个项目里面有详细的使用说明,很容易就上手了,非常感谢作者👍
  网络模型有了,加上数据就可以炼丹了,这里用的数据集是SeaShips(7000),格式是VOC格式,代码需要的也是VOC格式,所以改好文件夹路径就可以直接上了。训练用的是Google Colab,白嫖虽爽,但是也有很多不方便,后面再单独整理一篇使用记录吧

二、openCV-python

  openCV是很常用的图像处理库,支持python、C++等语言,这里用的是openCV-python
  openCV中的dnn模块实现了基于Darknet框架的YOLO v3网络,也就是说,模块给你搭好了网络,只需要你自己填入网络参数就好了,这个网络参数就是训练模型得到的权重文件
  训练模型时用的框架是Pytorch,而Pytorch框架下的权重文件是.pth格式,而openCV的dnn模块对于YOLO网络仅支持Darknet框架,于是乎,需要把.pth转为.weights格式(.weights是Darknet框架下的权重文件格式)

三、.pth转.weights

图1-.pth权重文件的网络结构示意图
  图1是训练得到的权重文件的网络结构,实际上这个.pth文件保存的是网络结构+权重参数,具体来说是嵌套的ordereddict
图2-转换格式的部分代码
  图2是转换格式的部分代码,因为在Darknet框架下,YOLO v3的cfg文件是不变的,因此对照cfg文件的网络配置,就可以一层一层的,按照顺序将.pth格式的网络权重参数(numpy数组)转换成.weights所需的二进制格式。如果需要转换代码的话,在评论区留邮箱吧

四 模型部署

  模型部署可以参考这篇博客的正文,代码都有注释,很容易理解


总结

  上面主要讲了对于YOLO v3网络,如何将.pth格式的权重文件转成.weights,其他的像训练模型、部署模型有很多现成的资料供参考,这里就不重复了

# -*- coding: gbk -*- import cv2 import numpy as np import torch from segment_anything import sam_model_registry, SamPredictor import os # 注册并加载 SAM 模型 sam_checkpoint = "/home/lichuang/project/Opencv-main/sam_yolo/yolov8_sam-weights/sam_vit_b_01ec64.pth" # 请根据实际路径修改 model_type = "vit_b" device = "cuda" if torch.cuda.is_available() else "cpu" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) predictor = SamPredictor(sam) # 定义输入和输出文件夹路径 input_root_folder = "/home/lichuang/project/Opencv-main/sam_yolo/traditional3d-plan1/dataset/input_images/" # 包含子文件夹的总文件夹路径 output_folder = "/home/lichuang/project/Opencv-main/sam_yolo/traditional3d-plan1/dataset/mask1/" # 生mask图片的输出文件夹路径 # 如果输出文件夹不存在,则创建它 if not os.path.exists(output_folder): os.makedirs(output_folder) # 遍历总文件夹下的所有子文件夹 for subfolder in os.listdir(input_root_folder): subfolder_path = os.path.join(input_root_folder, subfolder) if os.path.isdir(subfolder_path): # 在输出文件夹下创建对应的子文件夹 output_subfolder = os.path.join(output_folder, subfolder) if not os.path.exists(output_subfolder): os.makedirs(output_subfolder) # 遍历子文件夹中的所有图片 for filename in os.listdir(subfolder_path): if filename.endswith(('.png', '.jpg', '.jpeg')): img_path = os.path.join(subfolder_path, filename) # 读取图片 image = cv2.imread(img_path) if image is None: print(f"Failed to read {img_path}") continue image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 设置图像到预测器 predictor.set_image(image_rgb) # 简单地使用图像中心作为提示点(可根据需求优化提示点策略) height, width = image.shape[:2] input_point = np.array([[width // 2, height // 2]]) input_label = np.array([1]) # 进行预测 masks, scores, logits = predictor.predict( point_coords=input_point, point_labels=input_label, multimask_output=True, ) # 选择得分最高的 mask best_mask = masks[np.argmax(scores)] # 对掩码进行形态学操作,去除小孔洞和噪点(可选步骤,可根据实际情况调整) kernel = np.ones((5, 5), np.uint8) best_mask = cv2.morphologyEx(best_mask.astype(np.uint8), cv2.MORPH_OPEN, kernel) best_mask = cv2.morphologyEx(best_mask, cv2.MORPH_CLOSE, kernel) # 保存生的二值化mask图片 output_path = os.path.join(output_subfolder, filename) cv2.imwrite(output_path, best_mask * 255) print("Mask图片生!")此为掩码图的生代码,请修改代码,尝试解决此问题
最新发布
03-21
评论 102
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值