Para ejecutar YOLOv5, primero debe instalar el entorno de aprendizaje profundo. Para ver el tutorial, consulte Instalación del entorno de aprendizaje profundo de pytorch (versión GPU) .
El código de YOLOv5 es GitHub de código abierto - ultralytics/yolov5 en GitHub . Usar su código para satisfacer sus propias necesidades de detección de objetivos requiere tres pasos: 1. Preparar el conjunto de datos; 2. Configurar los parámetros del código y entrenar el modelo; 3. Predecir . El siguiente autor lo guiará paso a paso para realizar su propio entrenamiento de modelo de detección de objetivos.
1. Preparar el conjunto de datos
1.1 Recopilar imágenes
Recopilamos imágenes relevantes de acuerdo con nuestras propias necesidades, y aquí tomamos como ejemplo el reconocimiento de máscara. Recopilamos algunas imágenes de Internet con y sin máscaras, como se muestra a continuación: 
1.2 Use el software labelimg para etiquetar las imágenes recopiladas
1.2.1 Instalación del software labelimg
El software labelimg es una herramienta de etiquetado de datos de código abierto que puede etiquetar tres formatos. ① Archivo XML en formato de etiqueta VOC . ② archivo txt en formato de etiqueta yolo. ③crea un archivo json con formato de etiqueta ML.
La instalación de labelimg es muy sencilla, abrimos cmd e ingresamos el siguiente comando:
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple1.2.2 Etiquetado con software labelimg
En primer lugar, también podríamos crear una carpeta llamada VOC2007 y crear una carpeta llamada JPEGImages para almacenar las imágenes que recopilamos que deben etiquetarse; luego, crear una carpeta llamada Anotaciones para almacenar los archivos de etiquetas etiquetados; txt llamado predefined_classes.txt para almacenar el nombre de la clase que se va a marcar. La estructura se muestra en la siguiente figura:
Lo que queremos lograr aquí es detectar si se usa una máscara, por lo que solo hay 2 categorías en el archivo predefined_classes.txt, como se muestra en la siguiente figura:
Luego, debemos abrir cmd en el directorio de VOC2007 (debe estar en este directorio) e ingresar el siguiente comando:
labelimg JPEGImages predefined_classes.txtEste comando significa usar el software labelimg para etiquetar las imágenes en la carpeta JPEGImages según las categorías en el archivo predefined_classes.txt.
La interfaz abierta se muestra en la siguiente figura, donde
Open Dir es la carpeta a elegir para almacenar imágenes, y aquí nuestro comando lo establece por defecto en la carpeta JPEGImages;
Change Save Dir es la carpeta para cambiar la etiqueta de almacenamiento, aquí tenemos por defecto la carpeta Anotaciones;
PascalVOC es para seleccionar el formato de la etiqueta. Como se mencionó anteriormente, hay tres tipos principales. Por lo general, elegimos el formato xml de PascalVOC, y el formato YOLO también está bien, y los dos se pueden convertir entre sí;
Create RectBox es generar una línea de posición cruzada etiquetada para marcar la imagen.
Después de seleccionar la posición de detección del objetivo, aparecerá un cuadro de selección de etiqueta, y podemos seleccionar la etiqueta correspondiente, como se muestra en la figura a continuación. Luego puede hacer clic en Siguiente imagen para marcar la siguiente imagen hasta que todas las imágenes estén marcadas.
Los dos formatos de etiqueta se muestran en la siguiente figura:
El formato xml de PascalVOC:
El formato txt de YOLO:
1.3 Transformación de formato de etiqueta y división de conjunto de entrenamiento y conjunto de verificación
1.3.1 Convierta etiquetas de formato xml a formato txt y divida el conjunto de entrenamiento (80 %) y el conjunto de verificación (20 %)
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
classes = ["unmask", "mask"]
TRAIN_RATIO = 80 %训练集的比例
def clear_hidden_files(path):
dir_list = os.listdir(path)
for i in dir_list:
abspath = os.path.join(os.path.abspath(path), i)
if os.path.isfile(abspath):
if i.startswith("._"):
os.remove(abspath)
else:
clear_hidden_files(abspath)
def convert(size, box):
dw = 1. / size[0]
dh = 1. / size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def convert_annotation(image_id):
in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id)
out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' % image_id, 'w')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_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))
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
in_file.close()
out_file.close()
wd = os.getcwd()
wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir) # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
path = os.path.join(image_dir, list_imgs[i])
if os.path.isfile(path):
image_path = image_dir + list_imgs[i]
voc_path = list_imgs[i]
(nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
(voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
annotation_name = nameWithoutExtention + '.xml'
annotation_path = os.path.join(annotation_dir, annotation_name)
label_name = nameWithoutExtention + '.txt'
label_path = os.path.join(yolo_labels_dir, label_name)
prob = random.randint(1, 100)
print("Probability: %d" % prob)
if (prob < TRAIN_RATIO): # train dataset
if os.path.exists(annotation_path):
train_file.write(image_path + '\n')
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_train_dir + voc_path)
copyfile(label_path, yolov5_labels_train_dir + label_name)
else: # test dataset
if os.path.exists(annotation_path):
test_file.write(image_path + '\n')
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_test_dir + voc_path)
copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()1.3.2 Convierta el formato de etiqueta txt a formato xml y luego use el método de 1.3.1 para dividir el conjunto de datos
from xml.dom.minidom import Document
import os
import cv2
# def makexml(txtPath, xmlPath, picPath): # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径
def makexml(picPath, txtPath, xmlPath): # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径
"""此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件
在自己的标注图片文件夹下建三个子文件夹，分别命名为picture、txt、xml
"""
dic = {'0': "unmask", # 创建字典用来对类型进行转换
'1': "mask", # 此处的字典要与自己的classes.txt文件中的类对应，且顺序要一致
}
files = os.listdir(txtPath)
for i, name in enumerate(files):
xmlBuilder = Document()
annotation = xmlBuilder.createElement("annotation") # 创建annotation标签
xmlBuilder.appendChild(annotation)
txtFile = open(txtPath + name)
txtList = txtFile.readlines()
img = cv2.imread(picPath + name[0:-4] + ".jpg")
Pheight, Pwidth, Pdepth = img.shape
folder = xmlBuilder.createElement("folder") # folder标签
foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
folder.appendChild(foldercontent)
annotation.appendChild(folder) # folder标签结束
filename = xmlBuilder.createElement("filename") # filename标签
filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
filename.appendChild(filenamecontent)
annotation.appendChild(filename) # filename标签结束
size = xmlBuilder.createElement("size") # size标签
width = xmlBuilder.createElement("width") # size子标签width
widthcontent = xmlBuilder.createTextNode(str(Pwidth))
width.appendChild(widthcontent)
size.appendChild(width) # size子标签width结束
height = xmlBuilder.createElement("height") # size子标签height
heightcontent = xmlBuilder.createTextNode(str(Pheight))
height.appendChild(heightcontent)
size.appendChild(height) # size子标签height结束
depth = xmlBuilder.createElement("depth") # size子标签depth
depthcontent = xmlBuilder.createTextNode(str(Pdepth))
depth.appendChild(depthcontent)
size.appendChild(depth) # size子标签depth结束
annotation.appendChild(size) # size标签结束
for j in txtList:
oneline = j.strip().split(" ")
object = xmlBuilder.createElement("object") # object 标签
picname = xmlBuilder.createElement("name") # name标签
namecontent = xmlBuilder.createTextNode(dic[oneline[0]])
picname.appendChild(namecontent)
object.appendChild(picname) # name标签结束
pose = xmlBuilder.createElement("pose") # pose标签
posecontent = xmlBuilder.createTextNode("Unspecified")
pose.appendChild(posecontent)
object.appendChild(pose) # pose标签结束
truncated = xmlBuilder.createElement("truncated") # truncated标签
truncatedContent = xmlBuilder.createTextNode("0")
truncated.appendChild(truncatedContent)
object.appendChild(truncated) # truncated标签结束
difficult = xmlBuilder.createElement("difficult") # difficult标签
difficultcontent = xmlBuilder.createTextNode("0")
difficult.appendChild(difficultcontent)
object.appendChild(difficult) # difficult标签结束
bndbox = xmlBuilder.createElement("bndbox") # bndbox标签
xmin = xmlBuilder.createElement("xmin") # xmin标签
mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
xminContent = xmlBuilder.createTextNode(str(mathData))
xmin.appendChild(xminContent)
bndbox.appendChild(xmin) # xmin标签结束
ymin = xmlBuilder.createElement("ymin") # ymin标签
mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
yminContent = xmlBuilder.createTextNode(str(mathData))
ymin.appendChild(yminContent)
bndbox.appendChild(ymin) # ymin标签结束
xmax = xmlBuilder.createElement("xmax") # xmax标签
mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
xmaxContent = xmlBuilder.createTextNode(str(mathData))
xmax.appendChild(xmaxContent)
bndbox.appendChild(xmax) # xmax标签结束
ymax = xmlBuilder.createElement("ymax") # ymax标签
mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
ymaxContent = xmlBuilder.createTextNode(str(mathData))
ymax.appendChild(ymaxContent)
bndbox.appendChild(ymax) # ymax标签结束
object.appendChild(bndbox) # bndbox标签结束
annotation.appendChild(object) # object标签结束
f = open(xmlPath + name[0:-4] + ".xml", 'w')
xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
f.close()
if __name__ == "__main__":
picPath = "VOCdevkit/VOC2007/JPEGImages/" # 图片所在文件夹路径，后面的/一定要带上
txtPath = "VOCdevkit/VOC2007/YOLO/" # txt所在文件夹路径，后面的/一定要带上
xmlPath = "VOCdevkit/VOC2007/Annotations/" # xml文件保存路径，后面的/一定要带上
makexml(picPath, txtPath, xmlPath)
Si la etiqueta está en formato txt, es necesario prestar atención a varias cuestiones durante el proceso de conversión:
1. De acuerdo con las últimas líneas del código, la nota txt debe almacenarse en la carpeta YOLO;
2. Las etiquetas convertidas a formato xml se almacenarán en la carpeta Anotaciones;
3. Si hay un error de .shape, verifique si hay un archivo class.txt en la carpeta YOLO y elimínelo.
Con eso, la parte de preparar el conjunto de datos está completa.
2. Configurar parámetros de código y entrenar modelos
2.1 Descargar código fuente
El código de YOLOv5 es GitHub de código abierto - ultralytics/yolov5 en GitHub , y podemos descargar el código fuente en el sitio web. Como se muestra en la figura a continuación, aquí hemos seleccionado la versión v6.0.
Descomprimimos el código yolov5 descargado y luego lo abrimos con pycharm.Después de abrir, todo el directorio del código es el siguiente:
2.2 Unirse al conjunto de datos
Coloque el conjunto de datos que preparó en la carpeta VOCdevkit en el directorio del proyecto y use el método de conversión de etiquetas y división de conjuntos de datos presentado en 1.3 para dividir el conjunto de datos, como se muestra en la siguiente figura:
2.3 Configurar parámetros de código
1. Seleccione el entorno de pytorch configurado en la esquina inferior derecha de pycharm. Si el entorno no se ha instalado, consulte el artículo anterior;
2. Instale las bibliotecas dependientes requeridas por yolov5. Abra el terminal de comando de pycharm e ingrese el siguiente comando, como se muestra en la figura:
pip install -r requirements.txt
3. Descargue el archivo de prepeso. Descarga del peso previo del sitio web , aquí usamos yolov5s.pt , descárguelo y colóquelo en el directorio del proyecto.
4. Modifique el archivo de configuración de datos.
① Copie el archivo VOC.yaml en la carpeta de datos de este directorio, asígnele el nombre mask.yaml y modifíquelo consultando la figura siguiente.
②Copie una copia del archivo yolov5s.yaml en la carpeta del modelo en este directorio, asígnele el nombre yolov5s_mask.yaml y modifíquelo consultando la figura a continuación.
5. Ajuste de parámetros del archivo de tren
Como se muestra en la figura a continuación, la línea 436 es para configurar el archivo de peso previo, las líneas 437 y 438 son para configurar el archivo de configuración de datos, y la línea 440 es para configurar el número de iteraciones (puede configurarlo usted mismo según su necesidades).
Después de configurar los parámetros anteriores, puede ejecutar el archivo train.py para el entrenamiento. Pero puede encontrar los siguientes problemas:
Pregunta uno:
Esto significa que la memoria virtual no es suficiente. Podemos cambiar el parámetro num_workers nw en la línea 117 a 0 modificando el archivo datasets.py en la ruta utils.
Pregunta dos:
Esto indica un desbordamiento de memoria GPU. Podemos resolverlo reduciendo el tamaño del lote y el tamaño del parámetro de los trabajadores.
Después de completar la configuración de parámetros anterior, podemos entrenar nuestros propios datos. Ejecute el archivo train.py, la columna Ejecutar como se muestra en la siguiente figura indica que el entrenamiento ha comenzado.
3. Pronóstico
Una vez completado el entrenamiento, el proyecto tendrá una carpeta de carreras/entrenamiento/exp, que contiene los datos de peso entrenados y otros archivos de parámetros, como se muestra en la siguiente figura:
Luego abrimos el archivo detect.py y modificamos los siguientes parámetros. La línea 269 es para configurar el archivo de peso entrenado, la línea 270 es para configurar la carpeta de imágenes o la imagen específica que queremos detectar o llamar a la cámara, 0 significa llamar a la cámara.
Una vez completada la configuración, simplemente ejecute el archivo detect.py y los resultados de la detección se guardarán en la carpeta runs\detect\exp.
Los resultados de la prueba son los siguientes:
De esta forma, hemos completado un proyecto experimental basado en YOLOv5 entrenando nuestro propio conjunto de datos de detección de objetivos. ¿Le sorprenden los resultados de la prueba? También puede sentir que el proceso es engorroso o poco claro. La práctica hace al maestro, si practicas más, podrás captar los puntos principales pronto. Bien, gracias a todos, si encuentran otros problemas, pueden comunicarse.