【计算机视觉】目标检测—yolov5自定义模型的训练以及加载

1.摘要

目标检测是计算机视觉主要应用方向之一。目标检测通常包括两方面的工作,首先是招到目标,然后就是识别目标。目标检测可以分为单物体检测和多物体检测。常用的目标检测方法分为两大流派:一步走(one_stage)算法:直接对输入的图像应用算法并输出类别和相应的定位,典型的算法有yolo,ssd;两步走(two_stage)算法:先产生候选区域,然后在进行CNN分类,代表的算法有R-CNN。
本文重点讲述yolov5如何应用在实际项目中,主要包括两大部分,第一部分:如何使用yolov5训练自定义模型;第二部分:选练好自己的模型,如何加载并解读模型。

2.yolov5训练自定义模型

首先,去yolov5官网(https://github.com/ultralytics/yolov5)进行下载官方文件,也可以通过克隆地址(git clone https://github.com/ultralytics/yolov5.git)进行获取,并且下载预训练权重文件,放到yolov5文件下。结果如图:
在这里插入图片描述
下载好所有的文件后,接下来,我们需要进行配置环境,yolov5要求python>=3.7,pytorch>=1.7,我们可以通过
pip3 install -r requirements.txt进行其他安装包的安装,完成这些安装后,我们可以通过运行yolov5底下的detect.py文件来查看,环境是否配置成功。如果成功,则可以对data/image下的图片文件完成检测。原始模型可以识别80类目标,如下:
在这里插入图片描述
通过以上,我们可以检测预训练模型里面的目标,比如:球,人,大巴车等等,那该如何训练我们自己的数据集呢?通过labelImg将图片标注好,标签格式为yolo格式标签。将标注好的数据放在datasets/dataset_new/images,标签放在datasets/dataset_new/labels,各自底下分别有train,test,val三个文件。image格式为jpg格式,labels为文本文档。
在这里插入图片描述
一个图一个txt标注文件,每行一个物体,每行数据格式:类别id、x_center y_center width height,xywh必须归一化(0-1),其中x_center、width除以图片宽度,y_center、height除以画面高度,类别id必须从0开始计数。如下图所示。
在这里插入图片描述
准备好datasets文件后,将yolov5同级目录。YOLO会自动将…/datasets/dataset_new/images/train/1.jpg中的/images/替换成/labels/以寻找它的标签,如…/datasets/dataset_new/labels/train/1.txt。
准备好数据集后,将yolov5/data/coco128.yaml复制一份,重新命名为coco_1.yaml,并修改里面的文件路径和类别名称、数目。
在这里插入图片描述
复制models下对应模型的yaml文件,重命名,并修改nc值。
在这里插入图片描述
数据准备完成,模型需要修改的参数也完成了修改,接下里我们可以进行训练了train.py。
我们可以对trian.py里面的参数(如下图)按照自己的想法进行修改,也可以选择默认值。
在这里插入图片描述
训练需要一些时间,训练完成后,将得到我们需要的pt文件。下一步我们将加载与解读我们的模型。

3.模型的加载与解读

使用torch.hub.load()加载我们训练好的模型,如下,通过torch.hub.load进行加载模型,将图片导入,得到的results。可以通过results.pandas().xyxy[0]解读,结果是个张量,前四列为目标框的左上角和右下角,confidence为置信度,class为类别编号,name为类别名称。

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Image
im = 'https://ultralytics.com/images/zidane.jpg'

# Inference
results = model(im)

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

下面,是我加载自己训练模型的演示结果:

import cv2
import torch
import time
import numpy as np

model = torch.hub.load('./yolov5', 'custom', path='./weights/yolov5n_1.pt',source='local')
model.conf = 0.4

cap = cv2.VideoCapture(0)

fps_time = time.time()


while True:
    ret,frame = cap.read()
    frame = cv2.flip(frame,1)
    img_cvt = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)

    results = model(img_cvt)
    # print(results.pandas().xyxy[0].to_numpy())# tensor-to-numpy
    results_ = results.pandas().xyxy[0].to_numpy()
    i = 0
    for box in results_:
        l,t,r,b = box[:4].astype('int')
        confidence = str(round(box[4]*100,2))+"%"
        cls_name = box[6]

        if cls_name == "person":
            i += 1
        cv2.rectangle(frame,(l,t),(r,b),(0,255,0),2)
        cv2.putText(frame,cls_name + "-" + confidence,(l,t),cv2.FONT_ITALIC,1,(255,0,0),2)
    cv2.putText(frame, "person:"+str(i), (10, 20), cv2.FONT_ITALIC, 1, (0, 0, 0), 2)

    now = time.time()
    fps_txt = 1/(now - fps_time)
    fps_time = now

    cv2.putText(frame,str(round(fps_txt,2)),(50,50),cv2.FONT_ITALIC,1,(0,255,0),2)

    cv2.imshow("result",frame)
    if cv2.waitKey(10) & 0xFF == ord("q"):
        break

cap.release()
cv2.destroyAllWindows()

在这里插入图片描述

本篇主要给大家讲解如何使用yolov5端到端这个模型,希望对大家有多帮助。有疑问,可以进行交流。

由于目标检测模型的复杂性和训练数据的需求量较大,以下代码仅提供 PyTorch 中目标检测模型的基本框架和步骤,具体细节需要根据实际情况进行调整和优化。 1. 导入必要的库和模块 ```python import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms ``` 2. 定义模型结构 ```python class DetectionModel(nn.Module): def __init__(self): super(DetectionModel, self).__init__() # 定义模型结构,可以参考 YOLO、Faster R-CNN 等经典目标检测模型 ... def forward(self, x): # 定义前向传播过程,即模型如何从输入数据中提取特征并输出检测结果 ... return output ``` 3. 加载数据集 ```python transform = transforms.Compose([transforms.ToTensor(), ...]) trainset = torchvision.datasets.COCO(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2) testset = torchvision.datasets.COCO(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2) ``` 4. 定义损失函数和优化器 ```python criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) ``` 5. 训练模型 ```python for epoch in range(10): running_loss = 0.0 for i, data in enumerate(trainloader, 0): inputs, labels = data optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 2000 == 1999: print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000)) running_loss = 0.0 ``` 6. 测试模型 ```python correct = 0 total = 0 with torch.no_grad(): for data in testloader: images, labels = data outputs = net(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('Accuracy of the network on the 10000 test images: %d %%' % ( 100 * correct / total)) ``` 以上代码仅提供了目标检测模型的基本框架和步骤,具体细节需要根据实际情况进行调整和优化。另外,建议在训练模型时使用 GPU 加速,以提高训练速度和效果。
评论 15
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值