Обучение нейросети для CV на искусственно созданном датасете

Авторы статьи: Александр Михайлов, Александр Токарев.

Введение

Бывают ситуации, когда возникают сложности со сбором датасета для обучения нейросети, например:

Существующие в открытом доступе изображения отсутствуют в нужном количестве
Доступные изображения слишком дороги для использования в проекте
Требуется изображение сцен и объектов, которые крайне редко встречаются на практике или труднодоступны
Требуется добавить изображений в отдельные классы для выравнивания датасета

В этих ситуациях можно использовать генерацию изображений для обучения с помощью ИИ. На сегодня, передовые модели способны создавать изображения очень близкие к атуральным. Сгенерированный датасет можно разметить и обучить на нем нейросеть. В данной статье мы разберем как сгенерировать и разметить датасет при помощи SOTA моделей, а также обучить на созданном датасете модель для инференса. Рассмотрим генерацию датасета при помощи Stable Diffusion, разметку полученных изображений при помощи SAM и обучение на полученных данных модели YOLOv8.

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

Выбор генеративной модели

Мы выбрали Stable Diffusion из-за её впечатляющей фотореалистичности и открытого
доступа. Давайте немного углубимся в детали этой модели.
Stable Diffusion — это глубокая генеративная нейронная сеть, которая превращает текстовые запросы в изображения. Она имеет открытый код и веса, что позволяет использовать её на большинстве устройств с графическим процессором объемом видеопамяти не менее 4 ГБ. Модель предлагает несколько вариантов, оптимизированных под разные объемы памяти, так что вы сможете подобрать подходящий вариант под своё оборудование. К тому же, доступна облачная версия (https://stablediffusionweb.com/).
Модель принимает запросы на английском языке, хотя в сети можно найти русифицированные версии. Бесплатно можно создать до 10 изображений в день, а за дополнительные изображения взимается плата — около 4 руб. за изображение (цены на июнь 2024 г.). Оплата возможна с карт иностранных банков или через аналогичные ресурсы. В России сервис работает без VPN, доступен как онлайн, так и через API. Также Stable Diffusion можно установить в Google Colab, где он будет работать бесплатно.

При генерации изображений с помощью Stable Diffusion у пользователей есть возможность настроить различные параметры для достижения нужного результата.
Доступные параметры зависят от режима генерации. Вот некоторые из ключевых настроек:

Количество итераций (Diffusion Steps)
Начальное значение случайного числа (Random Seed)
Коэффициент содержания (Content Weight)
Температура (Sampling Temperature)
Преобразование стиля (Style Transfer)
Интерполяция изображений (Image Interpolation)

Подробнее о параметрах модели и особенностях написания промптов читайте в нашей
статье "Введение в генерацию изображений с помощью Stable Diffusion". По умолчанию
запросы должны быть на английском языке, хотя можно использовать и другие языки, но
результаты могут быть менее предсказуемыми. Чтобы получить наилучший результат,
запрос должен быть максимально точным.

Генерация датасета

Для начала мы составили простой промпт «a bear in the wood top view photostyle», но
полученный результат был не релевантным:

Следующим шагом было использование онлайн программы Clip Interrogator, которая по имеющемуся изображению составляет текстовый промпт для Stable Diffusion (описание работы с ней – в нашей статье «Введение в генерацию изображений с помощью Stable Diffusion»). В качестве исходного изображения использовали 5 найденных в интернете фотографий медведя с нужного верхнего ракурса. Интересно, что Clip Interrogator для каждой из 5 фотографий выдавал разные промпты. Подчистив эти промпты и слегка их меняя мы стали подавать их в Stable Diffusion и получать приемлемые изображения. Но это оказалось довольно трудоемко и все равно не гарантировало 100% приемлемого сюжета, так как Stable Diffusion реагирует на любые изменения в промпте и при этом не всегда предсказуемо. Тогда появилось желание автоматизировать процесс генерации. Самый простой способ – использовать изменение параметра seed – начального состояния генератора. Это дает разнообразие изображений в рамках одного промпта и гарантирует стабильность сюжета. Однако нам требовалось большее разнообразие сюжетов и тогда использовали орисанную ниже технологию. Если рассмотреть промпт, выданный программой Clip Interrogator, то видно, что он состоит из нескольких типовых частей. Рассмотрим один из вариантов:

Так как у нас было 5 разных промптов, то получилось по 5 вариантов для каждой части промпта. Далее составили программу, которая случайным образом из каждого списка извлекала термин и добавляла его к общему промпту. Для частей «Вид сверху» и «Прочее» выбирали по 2 выражения в одном промпте. Вы можете задавать различные правила комбинирования и получать самые разные сочетания в промпте. Также мы использовали негативный промпт, который был постоянный, хотя при желании, его можно сделать варьируемым.

Поэкспериментировав с промптами нам удалось сгенерировать 100 картинок довольно хорошего качества для формирования датасета. Такое количество может показаться маленьким, однако, как показала практика, его достаточно чтобы дообучить YOLO для пилотной версии продукта. В результате нам удалось получить примерно вот таких мишек.

Иногда правда получались и вот такие изображения с тремя ушами или другими
мутациями:

Часть 2. Сегментация изображений для датасета

Модель Segment Anything (SAM), выпущенная компанией Meta AI (Meta Platforms, Inc., деятельность которой запрещена в России) в апреле 2023 года, совершила революцию в области сегментации изображений. Эта модель получает входное изображение и автоматически сегментирует все объекты на нем, создавая маски сегментации. Затем эти маски передаются другим моделям, например тем, которые используются для преобразования изображений, чтобы создать маски указанных объектов. Расширенные возможности SAM являются результатом её обучения на миллионах изображений и масок, собранных с помощью "механизма обработки данных" типа "модель в цикле". Исследователи использовали SAM и его данные для интерактивного аннотирования изображений и обновления модели. Этот цикл повторялся много раз, чтобы улучшить как модель, так и набор данных. Подробнее читайте в нашей статье «Погружение в SAM от Meta AI: Всё, что нужно знать о сегментации изображений»

Список задач, которые решает Segment Anything Model:

Быстрое и качественное вырезание объектов
Отделение фона от объекта
Разделение изображения на составляющие
Вырезать любой объект на изображении одним щелчком

Для нашей задачи мы использовали фреймворк Autodistill с интегрированной моделью GroundingSAM. Этот фреймворк работает с использованием базовых моделей, таких как Grounding DINO и GroundedSAM, которые обычно слишком медленные для использования в режиме видеопотока. Одним из преимуществ этих моделей является их способность к нулевому снимку (zero shot object detection), что позволяет им работать без предварительной настройки на пользовательских данных, основываясь только на текстовых описаниях. https://blog.roboflow.com/launch-auto-label/

Посмотрим как GroundedSAM сегментирует реальное и сгенерированное изрбражения.

Это реальное изображение

Это изображение сгенерировано Stable Diffusion

Модель GroundedSAM использует Grounding DINO в качестве детектора объектов с открытым набором для интеграции с моделью сегментации объектов SAM (подробнее о том зачем нам использовать детектор перед сегментацией с помощью SAM так же читайте в статье «Погружение в SAM от Meta AI: Всё, что нужно знать о сегментации изображений») После этого можно добавлять метки к каждой из интересующих вас масок на вашем изображении. В приведенном примере вы можете включить в свой набор данных только те метки масок, которые подтверждают, что модели GroundedSAM обнаружили медведей.

Результат сегментации:

Сегментация реального изображения

Сегментация сгенерированного изображения

Здесь можно заметить, что объект на реальном изображении, в данном случае бурый медведь, находится далеко. Но, несмотря на это, SAM справляется хорошо. Это еще раз подтверждает, что Segment Anything Model предлагает мощное и универсальное решение для сегментации объектов на изображениях, позволяющее улучшать ваши наборы данных с помощью масок сегментации.

Часть 3. Обучение YOLOv8 на созданном и размеченном датасете

Итак, теперь у нас есть датасет для обучения модели. Мы выбрали одну из самых известных и проверенных моделей – Yolov8. Чтобы обучение проходило быстрее, выбрали самую маленькую модель Yolov8n. Обучение длилось 210 эпох и показало неплохие результаты. Проверка работоспособности модели проводилась на настоящих фото медведей, также найденных в интернете.

После тестирования обученной модели получили результат:

Как видим результат довольно хорош. Разумеется были и осечки, где модель нашла не
всех медведей либо приняла за медведя другой объект, но учитывая что мы брали для
обучения небольшой сет картинок, а также использовали самую маленькую модель,
можно считать эксперимент удавшимся. Гипотеза успешно проверена.

Выводы

В случае необходимости можно использовать технологию искусственной генерации датасета, имея несколько оригинальных изображений.
Для генерации изображений можно использовать пакет StablrDiffusion по API или установив его на локальный компьютер.
Создать качественный запрос можно с помощью программы Clip Interrogator. Чтобы получить разнообразный датасет можно комбинировать части запроса для Stable Diffusion между собой. Также можно использовать настройку seed.
Для разметки можно использовать предобученную модель сегментации SAM.