YOLOV5的训练与部署

最新推荐文章于 2025-05-02 19:07:48 发布

无敌暴龙战士832

最新推荐文章于 2025-05-02 19:07:48 发布

阅读量2.2k

点赞数 25

文章标签： YOLO python

本文链接：https://blog.youkuaiyun.com/2403_87255456/article/details/145601428

版权

本文简述使用自己电脑本地训练YOLOV5模型，并部署在树莓派上

机端环境部署

一、python环境搭建

先检查自己电脑是否有python以及它的版本

win+R打开运行窗口输入CMD

然后在命令符里输入python

可以看到我的版本是3.10的，如果没有的可以去官网下载https://www.python.org/

接下来是创建虚拟环境（可以使用anaconda这里就不在描述）

输入python -m venv yolov5 （yolov5是虚拟环境名字）

然后激活环境yolov5\Scripts\activate

最后这样即为成功创建虚拟环境

安装CUDA与cuDNN库

下载CUDA

下载前需要确认自己的版本驱动

输入nvidia-smi查看CUDA版本号和驱动版本，我这里是12.2，驱动版本是537

下面是CUDA 工具包和 CUDA 次要版本兼容性所需的最低驱动程序版本

具体可前往官网查看https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#id5

我这里选择12.1的下载链接https://developer.nvidia.com/cuda-toolkit-archive

建议选择自定义安装

然后一直下一步即可

下载cuDNN

下载与CUDA版本对应的cuDNN版本可去文档查看https://developer.nvidia.com/rdp/cudnn-archive

我这里选择v8.9.7的

下载完成后，需要把cuDNN文件中的三个文件夹bin、include和lib复制到CUDA的安装目录

安装PyTorch

我选择的版本是2.5.1的，尽量版本不要太高，会出现兼容问题https://pytorch.org/get-started/previous-versions/

接下来验证安装

输入python后

import torch
torch.__version__

出现我这样及安装成功

下载YOLOV5

去官网下载https://github.com/ultralytics/yolov5我的是yolov5-7.0的

下载后解压在一个文件夹并在同级目录下创建data文件夹

在data里面创建train和valid两个文件夹

接着分别在这两个文件夹中创建images和labels两个文件夹

在images文件夹放入需要标注的照片（train：valid=9：1）

下载labelimg标注工具

进入下载地址：https://github.com/tzutalin/labelImg/releases，选择windows_v1.8.1.zip，下载并解压，直接双击exe文件即可打开。

勾选自动保存

点击verify Image下面这个，使它变成YOLO

然后打开文件夹

保存路径为

这些是快捷键

！！！注意valid下的图片也需要标注

训练前准备

复制coco128.yaml文件，粘贴并改名为mydata.yaml

然后打开mydata.yaml，修改

path：是你的data路径

train：是你的训练集照片路径

val：是你的验证集路径

names：是你在labelimg中标注的标签名称

打开在yolov5-master下的train.py文件

修改后，我这里照片尺寸选定320z

在打开detect.py文件，修改照片尺寸

修改后

在models文件夹中，打开yolov5s.yaml（打开什么的.yaml）取决于你上面修改的两个文件中的--weights,将nc：后的数字改为你标注的种类数量

打开requirements.txt文件，将Export下的注释去掉

在命令行中进入yolov5-master目录

安装依赖

pip install -r requirements.txt

要是你报出这样的错误

经过我的实测，该错误对后面并没有影响，所以无需在意

开始训练

安装完依赖后，就可以使用python train.py开始训练

出现以下错误，使用pip install pyyaml

出现以下错误，使用pip install tqdm

出现以下错误，使用pip install pandas

出现以下错误，使用pip install requests

出现以下错误，使用pip install numpy==1.20

如果觉得这里太慢，可以ctrl+鼠标左键进入游览器下载

然后就可以开始训练了，出现的警告不用管，那是因为PyTorch版本过高导致的

训练完后他会告诉你模型放在哪个文件

导出模型

进入export.pt文件

修改--imgsz，修改的要与你之前训练前的大小一致

然后输入命令python export.py --weights runs\train\exp5\weights\best.pt --include onnx

看见该提示及成功转换

板端环境部署

我所使用的树莓派4B安装的是ubuntu20的，如果是树莓派官方镜像内容大致不差

安装虚拟环境

1.安装 Python 3 和 `venv` 模块

sudo apt install python3 python3-venv python3-pip

2.创建虚拟环境

python3 -m venv yolov5

3.激活虚拟环境

source yolov5/bin/activate

安装依赖

1.安装opencv

pip install opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple

2.安装pytorch

pip install torch==2.0.0 torchvision==0.15.0 -i https://mirrors.aliyun.com/pypi/simple/

3.安装onnx

pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

4.安装onnxruntime

pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

开始识别

1.建立文件夹存放onnx文件

mkdir YOLOV5_jx

2.将onnx文件复制到YOLOV5_jx中

可以借助vscode的远程连接，以及其他工具，这里不再过多描述

3.编写程序

在bast.onnx同一级目录下创建.py文件

touch run_YOLOV5.py

编写程序sudo nano run_YOLOV5.py

import cv2
import numpy as np
import onnxruntime as ort
import time
 
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param: 
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return
    """
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )
 
def _make_grid( nx, ny):
        xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
        return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)
 
def cal_outputs(outs,nl,na,model_w,model_h,anchor_grid,stride):
    
    row_ind = 0
    grid = [np.zeros(1)] * nl
    for i in range(nl):
        h, w = int(model_w/ stride[i]), int(model_h / stride[i])
        length = int(na * h * w)
        if grid[i].shape[2:4] != (h, w):
            grid[i] = _make_grid(w, h)
 
        outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
            grid[i], (na, 1))) * int(stride[i])
        outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat(
            anchor_grid[i], h * w, axis=0)
        row_ind += length
    return outs
 
 
 
def post_process_opencv(outputs,model_h,model_w,img_h,img_w,thred_nms,thred_cond):
    conf = outputs[:,4].tolist()
    c_x = outputs[:,0]/model_w*img_w
    c_y = outputs[:,1]/model_h*img_h
    w  = outputs[:,2]/model_w*img_w
    h  = outputs[:,3]/model_h*img_h
    p_cls = outputs[:,5:]
    if len(p_cls.shape)==1:
        p_cls = np.expand_dims(p_cls,1)
    cls_id = np.argmax(p_cls,axis=1)
 
    p_x1 = np.expand_dims(c_x-w/2,-1)
    p_y1 = np.expand_dims(c_y-h/2,-1)
    p_x2 = np.expand_dims(c_x+w/2,-1)
    p_y2 = np.expand_dims(c_y+h/2,-1)
    areas = np.concatenate((p_x1,p_y1,p_x2,p_y2),axis=-1)
    
    areas = areas.tolist()
    ids = cv2.dnn.NMSBoxes(areas,conf,thred_cond,thred_nms)
    if len(ids)>0:
        return  np.array(areas)[ids],np.array(conf)[ids],cls_id[ids]
    else:
        return [],[],[]
def infer_img(img0,net,model_h,model_w,nl,na,stride,anchor_grid,thred_nms=0.4,thred_cond=0.5):
    # 图像预处理
    img = cv2.resize(img0, [model_w,model_h], interpolation=cv2.INTER_AREA)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
 
    # 模型推理
    outs = net.run(None, {net.get_inputs()[0].name: blob})[0].squeeze(axis=0)
 
    # 输出坐标矫正
    outs = cal_outputs(outs,nl,na,model_w,model_h,anchor_grid,stride)
 
    # 检测框计算
    img_h,img_w,_ = np.shape(img0)
    boxes,confs,ids = post_process_opencv(outs,model_h,model_w,img_h,img_w,thred_nms,thred_cond)
 
    return  boxes,confs,ids
 
 
 
 
if __name__ == "__main__":
 
    # 模型加载
    model_pb_path = "best.onnx" #这里为自己训练出模型的名字
    so = ort.SessionOptions()
    net = ort.InferenceSession(model_pb_path, so)
    
    # 标签字典 （以下为标注时的种类）
    dic_labels= {0:'laji'}
    
    # 模型参数 （导出模型时的尺寸）
    model_h = 320
    model_w = 320
    nl = 3
    na = 3
    stride=[8.,16.,32.]
    anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
    anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(nl, -1, 2)
    
    video = 0
    cap = cv2.VideoCapture(video)
    flag_det = False
    while True:
        success, img0 = cap.read()
        if success:
            
            if flag_det:
                t1 = time.time()
                det_boxes,scores,ids = infer_img(img0,net,model_h,model_w,nl,na,stride,anchor_grid,thred_nms=0.4,thred_cond=0.5)
                t2 = time.time()
            
                
                for box,score,id in zip(det_boxes,scores,ids):
                    label = '%s:%.2f'%(dic_labels[id],score)
            
                    plot_one_box(box.astype(np.int16), img0, color=(255,0,0), label=label, line_thickness=None)
                    
                str_FPS = "FPS: %.2f"%(1./(t2-t1))
                
                cv2.putText(img0,str_FPS,(50,50),cv2.FONT_HERSHEY_COMPLEX,1,(0,255,0),3)
                
            
            cv2.imshow("video",img0)
        key=cv2.waitKey(1) & 0xFF    
        if key == ord('q'):
        
            break
        elif key & 0xFF == ord('s'):
            flag_det = not flag_det
            print(flag_det)
            
    cap.release()