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])