Conjunto de datos VOC a conjunto de datos YOLO

Cree su propio conjunto de datos de VOC

En primer lugar, almacene las imágenes y los archivos .xml en el conjunto de datos VOC de acuerdo con el siguiente formato de carpeta.

Las carpetas de mi conjunto de datos en todos los niveles se muestran en la figura:

Generar archivo de nombre de conjunto de datos

El tren.txt almacena los nombres de todas las imágenes del conjunto de datos (nota, ¡ es el nombre sin el sufijo !)

Como se muestra en la figura, hay varios archivos .jpg en mi carpeta de imágenes, en esta carpeta de imágenes Cree un .txt archivo con el contenido como se muestra en la figura:

 DIR  *.*/B>train.txt 

Modifique el sufijo del archivo .txt a .bat, como se muestra en la figura:

Haga doble clic en el archivo .bat y verá que se genera un archivo train.txt,

abra el archivo train.txt y verá que el contenido son los nombres de todos los archivos en la carpeta de imágenes. .

Elimine las dos líneas de list.bat y train.txt, dejando solo el nombre de la imagen
Nota: Dado que solo el nombre de la imagen se puede mantener en el train.txt final sin el sufijo , es necesario poner ".jpg " en cada línea del archivo train.txt "Eliminar, puede eliminar manualmente, pero si el conjunto de datos es demasiado grande, puede usar la siguiente secuencia de comandos de Python para eliminar:

import os
filename = r"train.txt"
new_filename = r"train1.txt"
with open(filename,encoding="utf-8") as f1, open(new_filename,"w",encoding="utf-8") as f2:
    for line in f1:
        new_line = line[:-5]
        f2.write(new_line)
        f2.write('\n')
f1.close()
f2.close()

El efecto de ejecución se muestra en la figura:

Abra el tren1.txt recién generado y podrá ver que los sufijos se han eliminado.

En este momento, copie el contenido de tren1.txt a \data\ImageSets\Main\train.txt.

Archivo .xml a archivo .txt

El modelo yolo de archivo .xml a .txt proporciona oficialmente el código de conversión, y el método de uso se muestra a continuación:

Cree un archivo .py con el siguiente contenido:

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test'] 
classes = ['red', 'yellow','green', 'turn_left',  'turn_right', 'stop']   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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('C:/Users/Lenovo/Desktop/data/Annotations/%s.xml' % (image_id),encoding='utf-8')
    out_file = open('C:/Users/Lenovo/Desktop/data/labels/%s.txt' % (image_id), 'w',encoding='utf-8')
    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
        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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()
for image_set in sets:
    if not os.path.exists('C:/Users/Lenovo/Desktop/data/labels/'):
        os.makedirs('C:/Users/Lenovo/Desktop/data/labels/')
    image_ids = open('C:/Users/Lenovo/Desktop/data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('%s.txt' % (image_set), mode='w', encoding='utf-8')
    for image_id in image_ids:
        list_file.write('C:/Users/Lenovo/Desktop/data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

Notas sobre el uso del código de conversión:

1. Cambie la sexta línea a 'train', 'val', 'test'
2. Cambie el contenido de las clases de la séptima línea al nombre de la categoría en su propio conjunto de datos
3. Preste atención a la escritura de cada dirección, asegúrese de usar barras invertidas

Después de ejecutar este script, encontrará que se genera una nueva carpeta de etiquetas en los datos, lo que significa que la conversión de datos es exitosa

Mejore su propio conjunto de datos de YOLO

Finalmente, copie la carpeta de etiquetas en el conjunto de datos de yolo.