【深度学习】YOLOv8训练，交通灯目标检测

一、数据处理

import traceback
import xml.etree.ElementTree as ET
import os
import shutil
import random
import cv2
import numpy as np
from tqdm import tqdmdef convert_annotation_to_list(xml_filepath, size_width, size_height, classes):in_file = open(xml_filepath, encoding='UTF-8')tree = ET.parse(in_file)root = tree.getroot()# size = root.find('size')# size_width = int(size.find('width').text)# size_height = int(size.find('height').text)yolo_annotations = []# if size_width == 0 or size_height == 0:for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = [float(xmlbox.find('xmin').text),float(xmlbox.find('xmax').text),float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text)]# 标注越界修正if b[1] > size_width:b[1] = size_widthif b[3] > size_height:b[3] = size_heighttxt_data = [((b[0] + b[1]) / 2.0) / size_width, ((b[2] + b[3]) / 2.0) / size_height,(b[1] - b[0]) / size_width, (b[3] - b[2]) / size_height]# 标注越界修正if txt_data[0] > 1:txt_data[0] = 1if txt_data[1] > 1:txt_data[1] = 1if txt_data[2] > 1:txt_data[2] = 1if txt_data[3] > 1:txt_data[3] = 1yolo_annotations.append(f"{cls_id} {' '.join([str(round(a, 6)) for a in txt_data])}")in_file.close()return yolo_annotationsdef main():classes = ["red", "green", "yellow", "off"]root = r"/ssd/xiedong/lightyolov5"img_path_1 = os.path.join(root, "Traffic-Lights-Dataset-Domestic/JPEGImages")xml_path_1 = os.path.join(root, "Traffic-Lights-Dataset-Domestic/Annotations")img_path_2 = os.path.join(root, "Traffic-Lights-Dataset-Foreign/JPEGImages")xml_path_2 = os.path.join(root, "Traffic-Lights-Dataset-Foreign/Annotations")dst_yolo_root = os.path.join(root, "Traffic-Lights-Dataset-YOLO")dst_yolo_root_img = os.path.join(dst_yolo_root, "images")os.makedirs(dst_yolo_root_img, exist_ok=True)dst_yolo_root_txt = os.path.join(dst_yolo_root, "labels")os.makedirs(dst_yolo_root_txt, exist_ok=True)index = 0img_path_1_files = os.listdir(img_path_1)xml_path_1_files = os.listdir(xml_path_1)for img_id in tqdm(img_path_1_files):# 右边的.之前的部分xml_id = img_id.split(".")[0] + ".xml"if xml_id in xml_path_1_files:try:new_name = f"{index:06d}.jpg"img = cv2.imdecode(np.fromfile(os.path.join(img_path_1, img_id), dtype=np.uint8), 1)  # img是矩阵cv2.imwrite(os.path.join(dst_yolo_root_img, new_name), img)new_txt_name = f"{index:06d}.txt"yolo_annotations = convert_annotation_to_list(os.path.join(xml_path_1, img_id[:-4] + ".xml"),img.shape[1],img.shape[0],classes)with open(os.path.join(dst_yolo_root_txt, new_txt_name), 'w') as f:f.write('\n'.join(yolo_annotations))index += 1except:traceback.print_exc()img_path_1_files = os.listdir(img_path_2)xml_path_1_files = os.listdir(xml_path_2)for img_id in tqdm(img_path_1_files):# 右边的.之前的部分xml_id = img_id.split(".")[0] + ".xml"if xml_id in xml_path_1_files:try:new_name = f"{index:06d}.jpg"img = cv2.imdecode(np.fromfile(os.path.join(img_path_2, img_id), dtype=np.uint8), 1)  # img是矩阵cv2.imwrite(os.path.join(dst_yolo_root_img, new_name), img)new_txt_name = f"{index:06d}.txt"yolo_annotations = convert_annotation_to_list(os.path.join(xml_path_2, img_id[:-4] + ".xml"),img.shape[1],img.shape[0],classes)with open(os.path.join(dst_yolo_root_txt, new_txt_name), 'w') as f:f.write('\n'.join(yolo_annotations))index += 1except:traceback.print_exc()if __name__ == '__main__':main()

二、环境

conda create -n py310_yolo8 python=3.10 -yconda activate py310_yolo8conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidiapip install ultralytics

data.yaml

path: /ssd/xiedong/lightyolov5/Traffic-Lights-Dataset-YOLO/
train: images
val: images
test: # test images (optional)# Classes
names:0: 'red'1: 'green'2: 'yellow'3: 'off'

三、训练

教程：

https://docs.ultralytics.com/modes/train/#comet

新建训练代码文件train.py：

from ultralytics import YOLO# Load a model
model = YOLO("yolov8s.pt")  # load a pretrained model (recommended for training)# Train the model with 2 GPUs
results = model.train(data="data.yaml", epochs=100, imgsz=640, device=[0, 1, 2, 3], batch=128)

开启训练：

python -m torch.distributed.run --nproc_per_node 4 train.py

结果会存在这里：

训练截图：

数据分布：

训练完成：

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size99/100      8.53G     0.7946     0.3619     0.7988         13        640: 100%|██████████| 104/104 [00:15<00:00,  6.91it/s]Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 207/207 [00:31<00:00,  6.62it/s]all      13194      31308        0.9      0.632       0.74       0.54Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size100/100      8.54G      0.744     0.3433     0.7975          2        640: 100%|██████████| 104/104 [00:15<00:00,  6.88it/s]Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 207/207 [00:31<00:00,  6.58it/s]all      13194      31308      0.902      0.632      0.739       0.54100 epochs completed in 1.325 hours.
Optimizer stripped from runs/detect/train3/weights/last.pt, 22.5MB
Optimizer stripped from runs/detect/train3/weights/best.pt, 22.5MBValidating runs/detect/train3/weights/best.pt...
Ultralytics YOLOv8.2.19 🚀 Python-3.10.14 torch-2.1.2 CUDA:0 (NVIDIA A100-PCIE-40GB, 40339MiB)CUDA:1 (NVIDIA A100-PCIE-40GB, 40339MiB)CUDA:2 (NVIDIA A100-PCIE-40GB, 40339MiB)CUDA:3 (NVIDIA A100-PCIE-40GB, 40339MiB)
Model summary (fused): 168 layers, 11127132 parameters, 0 gradients, 28.4 GFLOPsClass     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 207/207 [00:37<00:00,  5.51it/s]all      13194      31308      0.902      0.632      0.739       0.54red      13194      13150      0.958      0.768      0.879      0.666green      13194      16030      0.961      0.727      0.866      0.655yellow      13194        830      0.901       0.58      0.677      0.469off      13194       1298      0.787      0.455      0.536      0.369
Speed: 0.0ms preprocess, 0.5ms inference, 0.0ms loss, 0.5ms postprocess per image
Results saved to runs/detect/train3

数据下载

链接：https://pan.baidu.com/s/1z2wpvnt40y-o0228rU7IIQ?pwd=buzx 
提取码：buzx 

帮助：

https://docs.qq.com/sheet/DUEdqZ2lmbmR6UVdU?tab=BB08J2