Yolov5 crea un conjunto de datos y entrena el modelo de detección de objetivos

Recopilación de conjuntos de datos

Utilice labelimg para anotaciones de datos

Muestre algunos a continuación 内联代码片.

pip install labelimg
labelimg

El primer paso es hacer clic para abrir el archivo , seleccionar el directorio de imágenes a anotar, luego hacer clic para cambiar el directorio de almacenamiento y seleccionar la ubicación para guardar la etiqueta anotada.

Luego seleccione el formato de la etiqueta. Depende de la forma El modelo lee los datos durante el entrenamiento.
Se recomienda activar el guardado automático en Ver.

tecla de acceso rápido
A	Imagen anterior
D	siguiente imagen
W.	Selección rápida de fotogramas

Formato de conjunto de datos

Tome yolov5 como ejemplo

conjuntos de datos

tren

imágenes

etiquetas

prueba

imágenes

etiquetas

válido

imágenes

etiquetas

Se recomienda dividir el conjunto de datos en tren: prueba: válido = 7: 2: 1

Modelo de entrenamiento

Descargue y configure el entorno yolov5

Debe realizarse bajo la versión pytorch gpu

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

Colocar el conjunto de datos

Simplemente coloque el conjunto de datos en el directorio yolo5 y luego decore el archivo .yaml

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

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

Nómbrelo data.yaml y colóquelo en el directorio yolo5\dataset.

Modelo de entrenamiento

Abra train.py y configure los siguientes cuatro parámetros

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 es la ruta del conjunto de datos, use el archivo data.yaml
–batch-size es el número de imágenes en cada lote de entrenamiento. Incrementar este valor aumentará la velocidad de entrenamiento y también conducirá a un mayor uso de memoria
–tiempos de entrenamiento de épocas

El modelo utiliza yolov5s.pt de forma predeterminada. Los parámetros de cada modelo previamente entrenado son los siguientes:

Modelo	tamaño (píxeles)	valor mapa 0.5:0.95	mapa valor 0.5	Velocidad CPU b1 (ms)	Velocidad V100 b1 (ms)	Velocidad V100 b32 (ms)	parámetros (M)	FLOP @640 (G)
yolv5n	640	28.0	45,7	45	6.3	0,6	1.9	4.5
yolv5s	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
yolv5l	640	49.0	67,3	430	10.1	2.7	46,5	109.1
yolv5x	640	50,7	68,9	766	12.1	4.8	86,7	205.7
yolv5n6	1280	36.0	54.4	153	8.1	2.1	3.2	4.6
yolv5s6	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
yolv5l6	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 -

Luego ejecuta train.py para entrenar el 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 sin concentrador, detección de imagen ú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=)