Yolov5 cria um conjunto de dados e treina o modelo de detecção de alvo

Coleta de conjunto de dados

Use labelimg para anotação de dados

Mostre alguns abaixo 内联代码片.

pip install labelimg
labelimg

O primeiro passo é clicar para abrir o arquivo , selecionar o diretório da imagem a ser anotada, clicar para alterar o diretório de salvamento e selecionar o local de salvamento da etiqueta anotada.

Em seguida, selecionar o formato da etiqueta. Depende da forma como o modelo lê os dados durante o treinamento.
Recomenda-se ativar o salvamento automático na Visualização.

tecla de atalho
A	Foto anterior
D	próxima foto
C	Seleção rápida de quadros

Formato do conjunto de dados

Tome yolov5 como exemplo

conjuntos de dados

trem

imagens

rótulos

teste

imagens

rótulos

válido

imagens

rótulos

Recomenda-se que o conjunto de dados seja dividido em trem: teste: válido = 7: 2: 1

Modelo de treinamento

Baixe e configure o ambiente yolov5

Precisa ser executado na versão pytorch gpu

git clone https://github.com/ultralytics/yolov5 
cd yolov5
pip install -r requirements.txt  # 下载所需包

Coloque o conjunto de dados

Basta colocar o conjunto de dados no diretório yolo5 e decorar o arquivo .yaml

train: ../train/images
val: ../valid/images
test: ../test/images

nc: 你的数据种类数
names: ['类别1的名称', '类别2的名称', '类别3的名称']

Nomeie-o como data.yaml e coloque-o no diretório yolo5\dataset.

Modelo de treinamento

Abra train.py e defina os quatro parâmetros a seguir

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--data', type=str, default=ROOT / 'datasets/data.yaml', help='dataset.yaml path')
    parser.add_argument('--batch-size', type=int, default=32, help='total batch size for all GPUs, -1 for autobatch')
    parser.add_argument('--epochs', type=int, default=20, help='total training epochs')

–data é o caminho do conjunto de dados, use o arquivo data.yaml
–batch-size é o número de imagens em cada lote de treinamento. Aumentar esse valor aumentará a velocidade de treinamento e também levará a um maior uso de memória –tempos de treinamento de
épocas

O modelo usa yolov5s.pt por padrão. Os parâmetros de cada modelo pré-treinado são os seguintes:

Modelo	tamanho (pixels)	valor mAP 0,5:0,95	mAP val 0,5	Velocidade CPU b1 (ms)	Velocidade V100 b1 (ms)	Velocidade V100 b32 (ms)	parâmetros (M)	FLOP @640 (B)
YOLOv5n	640	28,0	45,7	45	6.3	0,6	1,9	4,5
YOLOv5s	640	37,4	56,8	98	6.4	0,9	7.2	16,5
YOLOv5m	640	45,4	64,1	224	8.2	1.7	21.2	49,0
YOLOv5l	640	49,0	67,3	430	10.1	2.7	46,5	109.1
YOLOv5x	640	50,7	68,9	766	12.1	4.8	86,7	205,7
YOLOv5n6	1280	36,0	54,4	153	8.1	2.1	3.2	4.6
YOLOv5s6	1280	44,8	63,7	385	8.2	3.6	12.6	16,8
YOLOv5m6	1280	51.3	69,3	887	11.1	6.8	35,7	50,0
YOLOv5l6	1280	53,7	71,3	1784	15,8	10,5	76,8	111,4
YOLOv5x6 + [TTA][TTA]	1280 1536	55,0 55,8	72,7 72,7	3136 -	26,2 -	19,4 -	140,7 -	209,8 -

Em seguida, execute train.py para treinar o modelo

Usar modelo

pytorch.hub

import torch

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

# Images
img = 'https://ultralytics.com/images/zidane.jpg'numpy, list

# Inference
results = model(img)

# Results
results.print()

Uso de modelos não-hub, detecção de imagem única

from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import Profile, check_img_size, non_max_suppression, scale_boxes
from utils.torch_utils import select_device
from utils.plots import Annotator, colors
import cv2

import torch


def dtect_img(model, img_path, save_path, size=640):

    imgsz = (640, 640)
    bs = 1  # batch_size
    conf_thres = 0.25
    iou_thres = 0.45
    max_det = 1000
    classes = None
    agnostic_nms = True

    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size


    dataset = LoadImages(img_path, img_size=imgsz, stride=stride, auto=pt)

    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

    # 数据读取
    for path, im, im0s, vid_cap, s in dataset:
        with dt[0]:
            im = torch.from_numpy(im).to(model.device)
            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim

        # Inference
        with dt[1]:
            pred = model(im)
        # NMS
        with dt[2]:
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

        det = pred[0]
        annotator = Annotator(im0s, line_width=3, example=str(names))
        if len(det):
            # 图像标注区域等比缩放
            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0s.shape).round()

            for c in det[:, 5].unique():
                n = (det[:, 5] == c).sum()  # detections per class
                s += f"{
      
      n} {
      
      names[int(c)]}{
      
      's' * (n > 1)}, "  # add to string

            # 打印检测结果
            *xyxy, conf, cls = det.tolist()[0]
            name = names[int(c)]
            conf = f'{
      
      float(conf):.2f}'
            print(xyxy, "可信度：", conf, name)

            # 图像标注
            label = name + " " + conf
            annotator.box_label(xyxy, label, color=colors(c, True))

            # 保存图片
            img = annotator.result()

    return img

if __name__ == '__main__':
     model_path = '' # 模型路径
     img_path = '' # 检测图片路径
     save_path = '' # 保存路径
     
     device = ''
     device = select_device(device)

     model_detect = DetectMultiBackend(model_path, device=device)
     img = dtect_img(model_detect, model_path, save_path=)