Усовершенствование алгоритма YOLO для точного прогнозирования более 600 пользовательских классов

Всем привет! В этом руководстве мы покажем, как обучить алгоритм YOLO (You Only Look Once) для обнаружения объектов с помощью пользовательских классов, в частности, из списка из более чем 600 классов. Алгоритм YOLO известен своей скоростью и точностью, что делает его идеальным для задач обнаружения объектов в реальном времени. К концу этого руководства у вас будет рабочая модель, которая может прогнозировать пользовательские классы на основе большого набора данных; В этом уроке мы сосредоточимся на создании системы обнаружения продуктов питания, в которой мы будем обнаруживать тако, креветки и овощи по изображениям и видео.

Вот краткое описание шагов, которым мы будем следовать:

Собрать и подготовить набор данных
Настройте YOLO для пользовательских классов
Обучите модель YOLO
Протестируйте и оцените модель

И. Соберите и подготовьте набор данных

Набор данных.

Google Open Images Dataset v6 — это крупномасштабный и разнообразный набор данных для визуального распознавания, содержащий более 9 миллионов изображений с аннотациями. Эти аннотации охватывают более 600 категорий объектов и охватывают различные приложения, такие как обнаружение объектов, обнаружение визуальных отношений и сегментация экземпляров.

OIDv6 организован в различные подмножества: обучение, проверка и тестирование. Аннотации включают более 16 миллионов ограничивающих прямоугольников объектов, 3 миллиона сегментаций экземпляров и 4 миллиона визуальных взаимосвязей. Кроме того, набор данных иерархически структурирован с помощью семантической онтологии, включающей более 1700 классов, что позволяет лучше понять отношения между различными объектами.

Исследователи и разработчики обычно используют набор данных Google Open Images для обучения и сравнительного анализа моделей компьютерного зрения. Он стал популярным ресурсом для проектов машинного обучения в различных областях, таких как беспилотные автомобили, робототехника и автоматическая модерация контента.

Создание нашего набора данных.

Мы будем использовать пакет OIDv6. Это пакет Python, доступный в индексе пакетов Python (PyPI). Этот пакет представляет собой простой в использовании набор инструментов, предназначенный для помощи пользователям в работе с набором данных Google Open Images.

Пакет OIDv6 позволяет пользователям:

Загрузите изображения и аннотации для определенных классов или групп классов из набора данных Google Open Images.
Фильтруйте изображения по наличию или отсутствию определенных классов.
Преобразуйте формат набора данных для работы с популярными платформами глубокого обучения, такими как TensorFlow и PyTorch.
Проверяйте и визуализируйте набор данных и аннотации с помощью встроенных инструментов.

Чтобы установить пакет OIDv6, вы можете использовать следующую команду pip:

pip install oidv6

После установки вы можете получить доступ к его функциям через интерфейс командной строки или импортировав пакет в свой скрипт Python. Документация пакета и примеры содержат дополнительные рекомендации по использованию инструментария для управления набором данных Google Open Images и управления им в соответствии с вашими требованиями.

Процесс загрузки набора данных

Первоначально мы должны создать текстовый файл под названием «classes.txt», в котором должны быть задокументированы обозначения интересующих классов. Для целей этого учебного руководства мы сосредоточимся на трех отдельных классах: тако, креветках и овощах.
Создав файл .txt, выполните следующие команды, чтобы загрузить набор данных, который будет использоваться для обучения нашей модели (предоставленная команда запускает инструмент загрузки Open Images Dataset V6 (OIDv6), который предназначен для извлечения и загрузки определенных изображений из комплексный набор данных Open Images. Команда предоставляет следующие параметры: en: означает, что язык набора данных — английский. type_data train: указывает, что данные Загружаемый тип предназначен для учебных целей. в данном случае тако, креветки и овощи). limit 100: это ограничение указывает, что следует загружать не более 100 изображений для каждого класса. multi_classes: этот флаг разрешает загрузчик для одновременного получения изображений, принадлежащих нескольким классам, вместо того, чтобы загружать изображения из каждого класса по отдельности.):

oidv6 downloader en - type_data train - classes ./classes.txt - limit 100 - multi_classes

3. Мы запускаем одну и ту же команду для проверки и тестирования наборов данных:

oidv6 downloader en - type_data validation - classes ./classes.txt - limit 50 - multi_classes
oidv6 downloader en - type_data test - classes ./classes.txt - limit 10 - multi_classes

II. Настройте YOLO для пользовательских классов

YOLO требует абсолютных маршрутов для обучения, поэтому мы должны разработать код для выполнения следующих задач:

Преобразуйте метки обнаружения объектов из формата Open Images Dataset V6 (OIDv6) в формат YOLO (You Only Look Once).
Создавайте списки файлов для обучения, проверки и тестовых данных.
Создайте объектный файл для YOLO, содержащий необходимые сведения о конфигурации.

Код начинается с импорта необходимых библиотек и определения глобальных констант и переменных. Затем он определяет несколько функций, в том числе print_msg для условной печати, get_classes для извлечения имен классов из файла classes.txt и label_contents для преобразования меток OIDv6 в формат YOLO.

Основная часть скрипта разделена на три раздела:

Перевести метки в формат YOLO. В этом разделе перебираются каталоги для данных обучения, проверки и тестирования, и для каждого файла изображения создается соответствующий файл этикетки в формате YOLO. Сценарий также создает текстовый файл для каждого каталога, содержащий список файлов изображений.
Создать списки файлов. В этом разделе создается список файлов для каждого каталога (обучения, проверки и тестирования), содержащий пути ко всем файлам изображений в соответствующем каталоге.
Создать объектный файл. В этом разделе создается файл «classes.txt» и файл «obj.data», содержащие количество классов, пути к спискам файлов обучения и проверки, а также файл «classes.txt». ” расположение файла. Он также указывает каталог резервного копирования.

from os import chdir, path, listdir, getcwd
import shutil
import cv2
import sys

DIRS = ["train", "validation", "test"]
DEBUG = True

SKIP_TRANSLATE_LABELS = False
SKIP_GENERATE_FILE_LISTS = False
SKIP_GENERATE_OBJ_FILE = False

if len(sys.argv) < 2:
    print("Missing classes file as argument")
    raise SystemExit

classes_file = path.realpath(sys.argv[1])

def print_msg(msg, isDebug=False):
    if not isDebug:
        print(msg)
    elif isDebug and DEBUG:
        print(msg)

def get_classes(classes_file):
    with open(classes_file) as f:
        return [l.strip().lower().replace(" ", "_") for l in f.readlines()]

def label_contents(img_filename, classes):
    # It assumes is already in the img_filename directory and that label file
    # is on "labels" directory
    img = cv2.imread(img_filename)
    height, width, _ = img.shape
    file_class = "_".join(path.basename(img_filename).split("_")[:-1])
    class_idx = classes.index(file_class)
    # OIDv6 Label data
    label_file = img_filename[:-4] + ".txt"
    label_lines = [line.strip() for line in open("labels/" + label_file, "r")]
    new_lines = []
    for label_line in label_lines:
        label, x1, y1, x2, y2 = label_line.split()
        x1, y1, x2, y2 = float(x1), float(y1), float(x2), float(y2)
        box_width = x2 - x1
        box_height = y2 - y1
        center_x = x1 + (box_width / 2.0)
        center_y = y1 + (box_height / 2.0)
        relative_cx = center_x / width
        relative_cy = center_y / height
        relative_bw = box_width / width
        relative_bh = box_height / height
        new_lines.append('{0} {1} {2} {3} {4}'.format(
            class_idx, relative_cx, relative_cy, relative_bw, relative_bh))
    return "\n".join(new_lines)


chdir(path.join("OIDv6", "multidata"))

######## Translate Labels to YOLO format  ########

if not SKIP_TRANSLATE_LABELS:
    classes = get_classes(classes_file)

    for DIR in DIRS:
        chdir(DIR)
        image_files = []
        for filename in listdir():
            labels_file = path.join(getcwd(), filename[:-4] + ".txt")
            if (    path.isfile(filename)
                    and filename.endswith(".jpg")
                    and not path.isfile(labels_file) ):
                image_files.append(filename)
                labels_file_contents = label_contents(filename, classes)
                display_filename = DIR + "/" + path.basename(labels_file)
                print_msg("Generating Labels File " + display_filename)
                with open(labels_file, "w") as f:
                    f.write(labels_file_contents + "\n")
        class_list_file = path.join("..", DIR + ".txt")
        with open(class_list_file, "w") as f:
            for image in image_files:
                f.write(f"{DIR}/{image}\n")
        chdir("..")
    print_msg("\n\n================= Label Translation Finished =================\n\n")

if not SKIP_GENERATE_FILE_LISTS:
    for DIR in DIRS:
        chdir(DIR)
        file_list = path.join("..", DIR + ".txt")
        with open(file_list, "w") as f:
            for filename in listdir(getcwd()):
                if filename.endswith(".jpg"):
                    f.write(path.join(DIR, filename) + "\n")
        print_msg(f"File List {DIR}.txt generated")
        chdir("..")
    print_msg("\n\n================= File Lists Generation Finished =================\n\n")

if not SKIP_GENERATE_OBJ_FILE:
    chdir("..")
    classes_file = shutil.copy(classes_file, path.join(getcwd(), "classes.txt"))
    num_classes = sum(1 for line in open(classes_file))
    with open(path.join(getcwd(), "obj.data"), "w") as f:
        f.write(f"classes={num_classes}\n")
        f.write(f"train=multidata/train.txt\n")
        f.write(f"valid=multidata/validation.txt\n")
        f.write(f"names={classes_file}\n")
        f.write(f"backup=./\n")
    print_msg("\n\n================= Object File Generation Finished =================\n\n")

Давайте создадим файл с именем yolo_preprocess_data.py на уровне сгенерированной папки OIDV и запустим предыдущий код, используя class.txt в качестве аргумента:

python yolo_preprocess_data.py ./classes.txt

Когда скрипт завершит работу, мы увидим это сообщение:

================= Создание объектного файла завершено =================

В OIDv6 у нас будут следующие файлы:

class.txt, multidata (папка) и obj.data (измените имя на objects.txt), заархивируйте эти файлы вместе и назовите заархивированный файл data.zip.

III. Обучите модель YOLO

Вам придется использовать графический процессор, в моем случае я использую Google Colab и Google Drive для загрузки data.zip.

Запустите следующие блоки кода:

A. Определения переменных ENV (определить модель для запуска, в моем случае я буду использовать yolov4-tiny и создать папку Training/Tacos на моем диске Google)

GOOGLE_COLAB_ENV = True
BACKUP_DIR = "Training/Tacos" # Make sure that your backup Directory exists
MODEL_TO_TRAIN = "yolov4-tiny" # (Only supported options: yolov4 or yolov4-tiny)

G_DRIVE_MOUNTPOINT = "/drive"
G_DRIVE_ROOT = G_DRIVE_MOUNTPOINT + "/MyDrive"
G_DRIVE_DATASETZIP = G_DRIVE_ROOT + "/Training/Data/dataset.zip"

from os import path, getcwd
if GOOGLE_COLAB_ENV:
    CONTENT = "/content"
    DATASET = CONTENT + "/multidata"
    SCRIPTS = CONTENT + "/yolov4-training-with-oidv6"
    DARKNET = CONTENT + "/darknet"
    BACKUP_DIR = G_DRIVE_ROOT + "/" + BACKUP_DIR
else:
    CONTENT = path.realpath(getcwd())
    DATASET = CONTENT + "/multidata"
    SCRIPTS = CONTENT
    DARKNET = CONTENT + "/../darknet"
    BACKUP_DIR = CONTENT + "/" + BACKUP_DIR

CFG_FILE = ""
PRE_TRAINED_WEIGHTS = ""
PTW_FILENAME = ""
CUSTOM_CFG_FILE = SCRIPTS + "/my-" + MODEL_TO_TRAIN + ".cfg"

if MODEL_TO_TRAIN == "yolov4":
    CFG_FILE = DARKNET + "/cfg/yolov4-custom.cfg"
    PRE_TRAINED_WEIGHTS_URL = "https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137"
    PTW_FILENAME = "yolov4.conv.137"
elif MODEL_TO_TRAIN == "yolov4-tiny":
    CFG_FILE = DARKNET + "/cfg/yolov4-tiny-custom.cfg"
    PRE_TRAINED_WEIGHTS_URL = "https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.conv.29"
    PTW_FILENAME = "yolov4-tiny.conv.29"

# You may edit this to point to your last weights to resume training
# PRE_TRAINED_WEIGHTS = "$BACKUP_DIR/my-yolov4-tiny_last.weights"
PRE_TRAINED_WEIGHTS = DARKNET + "/" + PTW_FILENAME

Б. Предоставление данных

# Only Colab's
from google.colab import drive
drive.mount(G_DRIVE_MOUNTPOINT)

!unzip '/drive/MyDrive/Colab Notebooks/yolo_oid_data/tacos/data.zip' -d "$CONTENT"
print("Dataset unzipped into " + CONTENT)

!mkdir -p "$BACKUP_DIR"
print("Backup Directory created at " + BACKUP_DIR)

C. Импорт и настройка зависимостей даркнета

!git clone --depth 1 https://github.com/AlexeyAB/darknet
%cd darknet
!sed -i 's/OPENCV=0/OPENCV=1/' Makefile
!sed -i 's/GPU=0/GPU=1/' Makefile
!sed -i 's/CUDNN=0/CUDNN=1/' Makefile
!sed -i 's/CUDNN_HALF=0/CUDNN_HALF=1/' Makefile
!make

D. КОНФИГУРАЦИЯ ФАЙЛА OBJECTS.TXT

# Set current valid absolute paths for the dataset information
escaped_content = (CONTENT + "/").replace("/", "\/")
escaped_bdir = BACKUP_DIR.replace("/", "\/")

!sed -i "s/train=/train=$escaped_content/" "$CONTENT"/objects.txt
!sed -i "s/valid=/valid=$escaped_content/" "$CONTENT"/objects.txt
!sed -i "s/names=/names=$escaped_content/" "$CONTENT"/objects.txt
!sed -i "s/backup=\.\//backup=$escaped_bdir/" "$CONTENT"/objects.txt

# Change to absolute paths the train/validation/test file lists
escaped_dataset = DATASET.replace("/", "\/")

!sed -i "s/^train/$escaped_dataset\/train/g" "$DATASET"/train.txt
!sed -i "s/^validation/$escaped_dataset\/validation/g" "$DATASET"/validation.txt
!sed -i "s/^test/$escaped_dataset\/test/g" "$DATASET"/test.txt

E. ФАЙЛЫ .CFG

Чтобы обучить модель YOLO Darknet, необходимо приобрести файл конфигурации (.cfg) и файл .weights. Файл .cfg необходимо настроить для размещения определенного количества классов, предназначенных для прогнозирования. Измените либо darknet/cfg/yolov4-tiny-custom.cfg, либо darknet/cfg/yolov4-custom.cfg в соответствии с вашими предпочтениями, а затем переименуйте файлы в my-yolov4-tiny.cfg или my-yolov4.cfg соответственно. . Чтобы реализовать эти изменения, вы можете создать автоматический скрипт или выполнить изменения вручную (в этом руководстве я буду использовать yolov4-tiny-custom.cfg):

[net]
max_batches = (# of Classes * 2000)
steps = (80% of max_batches), (90% of max_batches)
#### Last section of the configuration file
###### Three pairs if yolov4-custom, two pairs if yolov4-tiny-custom
[convolutional]
filters = ( (# of Classes + 5) * 3 )
[yolo]
classes = # of Classes
[convolutional]
filters = ( (# of Classes + 5) * 3 )
[yolo]
classes = # of Classes
[convolutional]
filters = ( (# of Classes + 5) * 3 )
[yolo]
classes = # of Classes

Загрузите измененный файл .cfg в свою среду и выполните следующий код:

CUSTOM_CFG_FILE = '/drive/MyDrive/Colab Notebooks/yolo_oid_data/Tacos/yolov4-tiny-custom.cfg'

Давайте проверим, что у нас есть все файлы, которые нам нужны:

print(f"./darknet detector train {CONTENT}/objects.txt {CUSTOM_CFG_FILE} {PRE_TRAINED_WEIGHTS} -dont-show -mjpeg_port 8090 -map")

F. Модель поезда

Выполните следующий код, чтобы начать процесс обучения модели. Обратите внимание, что продолжительность этой процедуры зависит от конкретной используемой модели и объема задействованных данных:

!./darknet detector train \
  "$CONTENT"/objects.txt \
  "$CUSTOM_CFG_FILE" \
  "$PRE_TRAINED_WEIGHTS" \
  -dont_show \
  -map

IV. Протестируйте и оцените модель

G. Тестовая модель

Чтобы протестировать модель, мы разработаем скрипт Python, который использует библиотеку OpenCV для обработки изображений и обнаружения объектов с использованием модели YOLOv4-tiny. Скрипт определяет три класса объектов: Taco, Shrimp и Vegetable.

Скрипт импортирует необходимые библиотеки: os, glob, cv2 и numpy.
Функция load_images принимает на вход path и возвращает список путей к файлам изображений с расширением .jpg.
Функция process_images обрабатывает список файлов изображений, используя прилагаемый файл конфигурации YOLOv4-tiny (cfg_file) и файл весов (weights_file). Он считывает каждое изображение и выполняет обнаружение объектов, используя модель YOLOv4-tiny. Обнаруженные объекты отображаются на изображении в виде ограничивающих прямоугольников вместе с именами их классов и оценками достоверности. Затем измененные изображения сохраняются в указанной выходной папке (output_folder).
Функция draw_boxes берет outputs из модели YOLOv4-tiny и исходное img. Он выполняет итерацию по обнаруженным объектам, проверяя, находятся ли их идентификаторы классов в списке желаемых идентификаторов классов и превышают ли их показатели достоверности 0,5. Если оба условия соблюдены, он вычисляет координаты ограничивающей рамки и рисует ограничивающую рамку на изображении с помощью функций OpenCV cv2.rectangle и cv2.putText.
Сценарий определяет пути к файлу конфигурации YOLOv4-tiny, файлу весов, папке входных изображений и выходной папке.
Наконец, сценарий вызывает функцию load_images для загрузки файлов изображений и функцию process_images для обнаружения объектов и сохранения обработанных изображений в указанной выходной папке.

Таким образом, этот сценарий выполняет обнаружение объектов на наборе изображений с использованием модели YOLOv4-tiny и сохраняет изображения с ограничивающими рамками и метками, указывающими обнаруженные объекты и их оценки достоверности.

import os
import glob
import cv2
import numpy as np

def load_images(path):
    image_files = glob.glob(os.path.join(path, '*.jpg'))
    return image_files

def process_images(image_files, cfg_file, weights_file, output_folder):
    net = cv2.dnn.readNetFromDarknet(cfg_file, weights_file)
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    
    for img_path in image_files:
        img = cv2.imread(img_path)
        blob = cv2.dnn.blobFromImage(img, 1/255.0, (416, 416), swapRB=True, crop=False)
        net.setInput(blob)
        
        layer_names = net.getLayerNames()
        output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers().flatten().tolist()]
        outputs = net.forward(output_layers)
        
        draw_boxes(outputs, img)
        
        output_file = os.path.join(output_folder, os.path.basename(img_path))
        cv2.imwrite(output_file, img)


def draw_boxes(outputs, img):
    class_ids = [0, 1, 2]  # Taco, Shrimp, Vegetable
    class_names = ['Taco', 'Shrimp', 'Vegetable']  # Add class names as strings

    for output in outputs:
        for detection in output:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            
            if class_id in class_ids and confidence > 0.5:
                center_x, center_y, w, h = (detection[0:4] * np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])).astype('int')
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                
                # Add confidence level to label and use class names instead of class ID
                label = f"{class_names[class_id]}: {confidence * 100:.2f}%"
                cv2.putText(img, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)


cfg_file = '/drive/MyDrive/Colab Notebooks/yolo_oid_data/Tacos/yolov4-tiny-custom.cfg'
weights_file = '/drive/MyDrive/Training/Tacos/yolov4-tiny-custom_best.weights'
images_path = '/content/multidata/test'
output_folder = '/drive/MyDrive/Training/Tacos/prediction2'
    
image_files = load_images(images_path)
process_images(image_files, cfg_file, weights_file, output_folder)

Результаты

Вы можете найти записную книжку для этого руководства и весь соответствующий код в следующем репозитории github:

https://github.com/agsmilinas/Training-YOLOv4-with-Custom-Dataset-from-Open-Images-Database-v6-OIDv6

Усовершенствование алгоритма YOLO для точного прогнозирования более 600 пользовательских классов

Вопросы по теме