Генерация label-файлов для обучения своей нейросетевой модели детекции с помощью библиотеки YOLO

Для обучения собственной модели детекции (обнаружения объектов на изображении, а точм числе их месторасположения), часто применяется следующий подход:
1) Выбирается необходимый для детекции объект (это может быть, например, определенное животное, модель автомобиля и другое, что необходимо выявлять на картинке);
2) Создается множество изображений с этим объектом (можно собственноручно произвести фотографирование объекта либо реализовать подбор картинок из интернета);
3) Множество изображений размечается, на картинке выделяется наименование и местоположение объекта (задаются координаты рамки, внутри которой объект находится). Эти данные сохраняются в специальный текстовый файл (будем называть его label-файлом), который подается скрипту обучения вместе с исходным изображением.

И вот в ситуации, когда изображений огромное количество (а зачастую чем больше датасет — тем обучение эффективнее), можно воспользоваться приемом для облегчения процесса генерации label-файлов. Суть его в следующем: мы берем уже готовую обученную модель, в данном примере это YOLO World, которая уже умеет осуществлять детекцию определенных классов объектов (на данный момент модель определяет порядка 80 классов, список можно посмотреть по ссылке: https://gist.github.com/rcland12/dc48e1963268ff98c8b2c4543e7a9be8), и если наш объект подпадает под какой-либо из перечисленных классов, то мы можем произвести автоматическую разметку и генерацию label-файлов «силами» готовой модели с минимальными трудозатратами.

Ссылка на блокнот с реализацией данного приема в Google Colab: https://colab.research.google.com/drive/1k3zW2DDiqAn114NHnDfaQyv2IshSAXC_?usp=sharing

Сначала устанавливаем модуль «ultralytics»:

# Устанавливаем ultralytics
!pip install ultralytics

Затем подключаем необходимые модули и также даем блокноту доступ к Google Диску:

# Подключаем необходимые модули
from ultralytics import YOLO
import os
import sys
# Подключаем Google Диск
from google.colab import drive
drive.mount('/content/drive')

Далее подключаем модель YOLO (в данном случае YOLO World, уже предобученную детекции по определенным классам объектов), а также устанавливаем ограничение по классу (выбираем из списка класс, наиболее близкий к распознаваемым объектам, для примера использовались изображения трамваев, которые в целом неплохо распознаются по классу поезда, что выяснилось опытным путем):

# Загрузка модели YOLOv8 World
model = YOLO('yolov8s-worldv2.pt')
# Ограничиваем детекцию выбранными классами, https://gist.github.com/rcland12/dc48e1963268ff98c8b2c4543e7a9be8
model.set_classes(['train'])

Последним этапом указываем пути к обрабатываемым изображениям и папку в которую необходимо сохранять label-файлы с результатами детекции:

# Конфиг для обработки
extensions = ('.jpg', '.JPG') # ограничиваем обрабатываемые файлы расширением
images_dir = '/content/drive/MyDrive/ml/trains/' # папка с изображениями
labels_dir = '/content/drive/MyDrive/ml/trains_labels/' # папка - куда сохранять метки

# Выполнение инференса для получения результатов и сохрание меток
for file in os.listdir(images_dir):
  image = os.fsdecode(file)
  if image.endswith(extensions):
    result = model(images_dir + image)
    result[0].save_txt(labels_dir + image.split('.')[0] + '.txt')

В результате в директории «labels_dir» должны сохраниться текстовые файлы с координатами объектов на фото, названия txt-файлов должны совпадать с названиями изображений (за исключением расширения), также стоит обратить внимание на выводимый лон работы скрипта, в случае не обнаружения объекта на фото — будет соответствующее уведомление, такие фото стоит перепроверить вручную.

Далее полученные label-файлы можно использовать для обучения уже своей модели!