Yolov7训练自己的数据集和ONNX/TRT部署

Yolov7训练自己的数据集和ONNX/Trt部署

一、环境配置

1.1 项目下载

项目原地址：GitHub - WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors

打开终端，输入以下命令将YOLOv7项目源码下载到本地：

git clone https://github.com/WongKinYiu/yolov7.git
# 进入yolov7的文件夹中
cd yolov7

1.2 环境配置

# 提取自己的项目所需的安装包和版本（针对于已经有环境的项目进行移植，如果第一次弄忽略该步骤）
pip freeze > requirements.txt
# pip批量安装requirements.txt文件中包含的组件依赖
pip install -r requirements.txt# conda批量导出包含环境中所有组件的requirements.txt文件
conda list -e > requirements.txt
# conda批量安装requirements.txt文件中包含的组件依赖
conda install --yes --file requirements.txt

【这里发给出我自己用的代码版本，也可以直接使用官网给出的。缺什么后面自己在配置】

1.3 Docker 使用yolov7【优点是稳定的运行环境，如果是修改yolov7框架的不建议使用】【只是给出官网教程】

# 设置docker环境，确保docker和nvidia-docker的安装，创建启动docker容器
# create the docker container, you can change the share memory size if you have more.
nvidia-docker run --name yolov7 -it -v your_coco_path/:/coco/ -v your_code_path/:/yolov7 --shm-size=64g nvcr.io/nvidia/pytorch:21.08-py3# apt install required packages
# 安装依赖 在Docker容器内部执行以下命令来安装必要的Python包：
apt update
apt install -y zip htop screen libgl1-mesa-glx
# pip install required packages
pip install seaborn thop# go to code folder
cd /yolov7

二、数据准备

2.1 数据集搜寻

🐞高效搜寻数据集的宝藏平台：遇见数据集整合搜索🐞

🔗 https://www.selectdataset.com/

为什么选择遇见数据集？—— 集成搜索的核心优势

一站式资源聚合
平台将 Kaggle、UCI、Google Dataset Search 等数十个主流数据集平台的资源整合于一体，无需在多个网站间切换，输入关键词即可同步获取多源结果，大幅提升检索效率。
智能关联与跳转
搜索结果直接展示各平台的数据集链接，支持一键跳转至原始页面。例如搜索 “医学影像数据集”，可同时呈现 Kaggle 的公开数据集、UCI 的医学数据库及相关学术机构的资源入口，避免重复筛选。
精准分类与筛选
提供领域标签（计算机视觉 / 自然语言处理 / 金融等）、数据格式（CSV/JSON/ 图像）、许可类型（开源 / 学术使用）等多维筛选条件，帮助用户快速定位符合需求的数据集。

2.2 数据标注 [lableimg和CVAT]

2.2.1 常用的Lableimg

使用命令行安装

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple# 安装完成之后输入下面命令即可启动
labelimg

安装命令行成功图：
在这里插入图片描述

view里面有个auto save mode 点击之后就可以自动保存标签
## 常用快捷键：
A：切换到上一张图片
D：切换到下一张图片
W：调出标注十字架
del ：删除标注框

在这里插入图片描述

标注报错问题解决

libpng warning: iCCP: known incorrect sRGB profile
libpng warning: iCCP: known incorrect sRGB profile
Traceback (most recent call last):File "C:\Users\Jorya\.conda\envs\yolov8\Lib\site-packages\libs\canvas.py", line 530, in paintEventp.drawLine(self.prev_point.x(), 0, self.prev_point.x(), self.pixmap.height())
TypeError: arguments did not match any overloaded call:drawLine(self, l: QLineF): argument 1 has unexpected type 'float'drawLine(self, line: QLine): argument 1 has unexpected type 'float'drawLine(self, x1: int, y1: int, x2: int, y2: int): argument 1 has unexpected type 'float'drawLine(self, p1: QPoint, p2: QPoint): argument 1 has unexpected type 'float'drawLine(self, p1: Union[QPointF, QPoint], p2: Union[QPointF, QPoint]): argument 1 has unexpected type 'float'

解决方案

这个错误是由于 PyQt5 的drawLine方法需要整数参数，而代码中传入了浮点数导致的。在 Python 3 中，除法运算默认返回浮点数，而 PyQt5 的绘图函数要求坐标点为整数类型。

# 你需要修改canvas.py文件中第 530 行的代码，将浮点数坐标转换为整数。
# 打开文件C:\Users\Jorya\.conda\envs\yolov8\Lib\site-packages\libs\canvas.py，找到第 530 行：
原代码（错误行）：
p.drawLine(self.prev_point.x(), 0, self.prev_point.x(), self.pixmap.height())
修改后：
p.drawLine(int(self.prev_point.x()), 0, int(self.prev_point.x()), self.pixmap.height())# 文件中第531  行的代码，将浮点数坐标转换为整数。
原代码（错误行）：
p.drawLine(0, self.prev_point.y(), self.pixmap.width(), self.prev_point.y())
修改后：
p.drawLine(0, int(self.prev_point.y()), self.pixmap.width(), int(self.prev_point.y()))# 文件中第 526 行的代码，将浮点数坐标转换为整数。
原代码（错误行）：
p.drawRect(left_top.x(), left_top.y(), rect_width, rect_height)
修改后：
p.drawRect(int(left_top.x()), int(left_top.y()), int(rect_width), int(rect_height))

标注完整之后会有两个主要的文件夹，一个是图片文件夹另外一个是标签文件夹

2.2.2 私有化部署和协同CVAT

官网地址

2.3 数据划分

针对图片和标签都在一个文件夹中

-*- coding: utf-8 -*-
import os
import random
import shutil
from tqdm import tqdm
import re
# 其中数据集和标签都在一个文件夹中，通过读取后缀进行替换
"""
最终划分成yolo可训练数据集格式:
|——images|-test|-train|-valid
|——labels|-test|-train|-valid
|——test.txt
|——train.txt
|——val.txt
## 遇到的yolov7训练问题，已经把预训练模型给下载下来了，进行train训练的时候。
## 一直提示git问题：subprocess.CalledProcessError: Command 'git tag' returned non-zero exit status 128.
## 解决方法：使用git拉取yolov7框架，在进行训练
# nohup python   train.py  --batch 48 --epochs 150 --cfg cfg/deploy/yolov7.yaml --weights yolov7.pt --device 0,1,2 > train.log 2>&1 &
"""
# 数据集文件夹路径
data_dir = "datasets/to/your/path"
# 创建一个空列表来存储有效图片和标签的路径
valid_images = []
valid_labels = []  
# 遍历 data 文件夹下的所有文件
for file_name in os.listdir(data_dir):# 获取文件的完整路径file_path = os.path.join(data_dir, file_name)# 获取文件的扩展名ext = os.path.splitext(file_name)[-1].lower()if ext.lower()  in ['.jpg', '.png' ,'.jpeg']:  # 处理图片文件# 根据扩展名替换成对应的 label 文件名label_name = file_name.replace(ext, ".txt")label_path = os.path.join(data_dir, label_name)# 判断 label 是否存在if not os.path.exists(label_path):# 删除图片os.remove(file_path)print("Deleted:", file_path)else:# 将图片路径添加到列表中valid_images.append(file_path)# 将 label 路径添加到列表中valid_labels.append(label_path)
# 确保每个标签都有对应的图片
valid_labels = [label for label in valid_labels if os.path.exists(label)]
# 创建目标目录结构
base_dir = "newdatasets/to/your/path"
for split in ["train", "valid", "test"]:os.makedirs(os.path.join(base_dir, "images", split), exist_ok=True)os.makedirs(os.path.join(base_dir, "labels", split), exist_ok=True)
# 打开三个txt文件用于存储分配的路径
train_txt = open(os.path.join(base_dir, "train.txt"), "w")
valid_txt = open(os.path.join(base_dir, "val.txt"), "w")
test_txt = open(os.path.join(base_dir, "test.txt"), "w")
# 初始化计数器
train_count = 0
valid_count = 0
test_count = 0
# 遍历每个有效图片路径
for i in tqdm(range(len(valid_images))):image_path = valid_images[i]label_path = re.sub(r'\.(jpg|jpeg|png)$', '.txt', image_path, flags=re.IGNORECASE)# 随机生成一个概率r = random.random()# 判断图片和标签应该移动到哪个文件夹 (train:valid:test = 7:2:1)if r < 0.1:# 移动到 test 文件夹destination = "test"image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))test_txt.write(f"{image_destination}\n")  # 写入到 test.txttest_count += 1elif r < 0.2:# 移动到 valid 文件夹destination = "valid"image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))valid_txt.write(f"{image_destination}\n")  # 写入到 valid.txtvalid_count += 1else:# 移动到 train 文件夹destination = "train"image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))train_txt.write(f"{image_destination}\n")  # 写入到 train.txttrain_count += 1# 生成目标文件夹中 images 和 labels 的路径image_destination_path = os.path.join(base_dir, "images", destination, os.path.basename(image_path))label_destination_path = os.path.join(base_dir, "labels", destination, os.path.basename(label_path))# 移动图片到目标文件夹shutil.move(image_path, image_destination_path)# 移动标签到目标文件夹if os.path.exists(label_path):shutil.move(label_path, label_destination_path)
# 关闭txt文件
train_txt.close()
valid_txt.close()
test_txt.close()
# 输出有效图片和标签的数量
print(f"有效图片数: {len(valid_images)}")
print(f"有效标签数: {len(valid_labels)}")
# 输出 train, valid, test 集合中的图片数量
print(f"训练集图片数: {train_count}")
print(f"验证集图片数: {valid_count}")
print(f"测试集图片数: {test_count}")

三、模型训练

3.1 创建yaml文件

创建一个hat.yaml文件在yolov7文件夹中的data文件夹中
在这里插入图片描述

3.2 模型训练代码和参数解释【单卡训练和多卡训练】

单卡训练：

# train p5 models
python train.py --workers 8 --device 0 --batch-size 32 --data data/hat.yaml --img 640 640 --cfg cfg/training/yolov7.yaml --weights '' --name yolov7 --hyp data/hyp.scratch.p5.yaml# train p6 models
python train_aux.py --workers 8 --device 0 --batch-size 16 --data data/hat.yaml --img 1280 1280 --cfg cfg/training/yolov7-w6.yaml --weights '' --name yolov7-w6 --hyp data/hyp.scratch.p6.yaml

多卡训练

# train p5 models
python -m torch.distributed.launch --nproc_per_node 4 --master_port 9527 train.py --workers 8 --device 0,1,2,3 --sync-bn --batch-size 128 --data data/hat.yaml --img 640 640 --cfg cfg/training/yolov7.yaml --weights '' --name yolov7 --hyp data/hyp.scratch.p5.yaml# train p6 models
python -m torch.distributed.launch --nproc_per_node 8 --master_port 9527 train_aux.py --workers 8 --device 0,1,2,3,4,5,6,7 --sync-bn --batch-size 128 --data data/hat.yaml --img 1280 1280 --cfg cfg/training/yolov7-w6.yaml --weights '' --name yolov7-w6 --hyp data/hyp.scratch.p6.yaml

3.3 模型验证

# 得到的模型进行单张验证
python detect.py --weights best.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg
# 得到的模型进行文件夹中的图片验证
python detect.py --weights best.pt --conf 0.25 --img-size 640 --source inference/images/

四、部署ONNX和TRT模式的代码

4.1 模型导出

wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
python export.py --weights ./yolov7-tiny.pt --grid --end2end --simplify --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 --img-size 640 640
git clone https://github.com/Linaom1214/tensorrt-python.git
python ./tensorrt-python/export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16

4.2 ONNX模型使用代码

import sys
import torch
print(f"Python version: {sys.version}, {sys.version_info} ")
print(f"Pytorch version: {torch.__version__} ")
# Inference for ONNX model
import cv2
cuda = True
w = "/content/yolov7/yolov7-tiny.onnx"
img = cv2.imread('/content/yolov7/inference/images/horses.jpg')
import cv2
import time
import requests
import random
import numpy as np
import onnxruntime as ort
from PIL import Image
from pathlib import Path
from collections import OrderedDict,namedtuple
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
session = ort.InferenceSession(w, providers=providers)
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):# Resize and pad image while meeting stride-multiple constraintsshape = im.shape[:2]  # current shape [height, width]if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# Scale ratio (new / old)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])if not scaleup:  # only scale down, do not scale up (for better val mAP)r = min(r, 1.0)# Compute paddingnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh paddingif auto:  # minimum rectangledw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh paddingdw /= 2  # divide padding into 2 sidesdh /= 2if shape[::-1] != new_unpad:  # resizeim = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add borderreturn im, r, (dw, dh)
names = ['hat']
colors = {name:[random.randint(0, 255) for _ in range(3)] for i,name in enumerate(names)}
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = img.copy()
image, ratio, dwdh = letterbox(image, auto=False)
image = image.transpose((2, 0, 1))
image = np.expand_dims(image, 0)
image = np.ascontiguousarray(image)
im = image.astype(np.float32)
im /= 255
im.shape
outname = [i.name for i in session.get_outputs()]
outname
inname = [i.name for i in session.get_inputs()]
inname
inp = {inname[0]:im}
# ONNX inference
outputs = session.run(outname, inp)[0]
outputs
ori_images = [img.copy()]
for i,(batch_id,x0,y0,x1,y1,cls_id,score) in enumerate(outputs):image = ori_images[int(batch_id)]box = np.array([x0,y0,x1,y1])box -= np.array(dwdh*2)box /= ratiobox = box.round().astype(np.int32).tolist()cls_id = int(cls_id)score = round(float(score),3)name = names[cls_id]color = colors[name]name += ' '+str(score)cv2.rectangle(image,box[:2],box[2:],color,2)cv2.putText(image,name,(box[0], box[1] - 2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225, 255, 255],thickness=2)  
Image.fromarray(ori_images[0])