По легенде, во время премьеры первого в истории человечества фильма «Прибытие поезда на вокзал Ла-Сьота» зрители устроили панику — они решили, что поезд настоящий и движется прямо на них. Больше века спустя на поезда смотрят нейросети — и даже принимают 3D-модели вагонов за реальные.

Я Никита Каменев, видеоаналитик СИБУР Диджитал. Сегодня расскажу, как мы учили нейросеть распознавать прибытие поезда на 3D-моделях из Unreal Engine.

Исходная задача

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

• Положения поезда — стоит, движется или эстакада пустая;

• Положения мостиков — подняты или опущены;

• Рабочих — нужно отличить наших сотрудников от рабочих со стороны подрядчика;

• Наличие и отсутствие привязи — так как это высотные работы, нужно следить за безопасностью, и сообщать оператору, если на рабочем нет страховки.

Разработку нейросети мы начали со сбора датасета – нам скинули изображения с камеры видеонаблюдения, на которых видно поезда и рабочих в разных положениях и состояниях. На сбор данных ушёл месяц, на обучение модели — ещё один. Потом тестирование и дообучение, в общем, здесь всё стандартно.

Модель в итоге оказалась настолько хороша, что привела к перевыполнению плана — и заказчик предложил масштабировать её на другие площадки.

Задача со звёздочкой

Изначально мы обучили классификатор на данных с одной конкретной эстакады, которая находилась в закрытом помещении. Что будет, если показать нейросети прибытие поезда на уличную эстакаду?

Можно предположить, что модель, обученная на закрытой эстакаде, так же хорошо покажет себя и на открытой, но нет. Она не учитывает дождь, снег, время суток, засветы от солнца, туман и другие факторы. В текущем виде выпускать на другие площадки её нельзя — нужно дообучение. Но на уличных площадках факторов много: с учётом всех сезонов и погодных условий реальные данные в идеале нужно собирать целый год. Неужели нельзя побыстрее?

Бистро для алгоритма

Расширить датасет без сбора настоящих новых данных можно тремя способами:

• Аугментация — модифицируем то, что уже есть.

• Open source — скачиваем чужие данные из интернета и размечаем их под себя.

• GAN — обучить нейросеть на данных, сгенерированных другой нейросетью.

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

В этом случае не справилась даже нейросеть — например, приняла двух человек за одного, просто потому что они стояли рядом:

Сначала мы пробовали решить эту проблему всё той же аугментацией. Брали сложные ракурсы и добавляли на них людей с других изображений, вырезая их в GIMP и фотошопе. Это сэкономило время — не нужно было ждать, когда мы получим те же сложные ракурсы с людьми. Но результаты выглядели нереалистично. Поэтому мы пошли другим путём.

Строим игровой завод

Так ли принципиально обучать нейросеть на настоящих кадрах и видео с камер, если алгоритм просто ищет контуры нужных объектов в разных состояниях? Нет — поэтому мы решили построить эстакаду в игровом движке Unreal Engine. Для генерации картинок взяли открытый сетпак с 3D-моделями разных модулей завода, построили пространство и пустили гулять ботов-рабочих. Вот как это выглядело:

После обучения на таких данных точность определения рабочих возросла с 81% до 92%. Да, всего на 11% лучше, но эта разница чувствуется в проде. Тем не менее этого всё равно мало — заводы в Unreal Engine выглядят игрушечными, что бы мы ни делали.

Вот бы был способ перенести реальные эстакады в пространство Unreal Engine… В теории, это можно сделать с помощью NeRF.

Игра на NeRFах

NeRF — это Neural Radiance Field, нейросеть, которая создаёт 3D-модели на основе видео и фотографий. Чтобы сделать игровой завод из Unreal Engine менее игрушечным, мы попробовали действовать так:

• Берём телефон и устраиваем прогулку по заводу.

• Создаём реалистичную 3D-модель с NeRF.

• Закидываем эту модель в Unreal Engine и на её основе создаём датасет.

Но обучение NeRF и генерация моделей — это отдельная задача. У нас был выбор из трёх инструментов: Luma AI, Nerfstudio и Instant NeRF.

В выборе инструмента важно учесть ещё и то, что нам все эти модели надо перенести в Unreal Engine. Так, Luma AI может при помощи отдельных плагинов экспортировать в платном формате .nvol с помощью плагина volinga или в бесплатном формате .luma. Nerfstudio тоже позволяет использовать volinga для экспорта. Но есть вариант лучше — можно экспортировать модель как облако точек для Blender и потом перенести её в движок. Та же схема применима и к Instant Nerf. 

В итоге мы выбрали экспорт через Blender — формат .luma не позволяет использовать промышленные данные, которые не надо светить в открытых источниках. Полученные таким способом модели выглядели как из снов:

Проблема моделей больших объектов в том, что результат сильно зависит от исходных данных. Мы эстакады обходили пешком и снимали видео на телефон с широкоугольным объективом. Из-за этого местами сетка ломалась и мы получали странные результаты с размытием и чёрными пятнами. В идеале для таких задач надо облететь объект на квадрокоптере. У нас такой возможности не было, поэтому мы дорабатывали модели из NeRF в движке.

Модели поездов мы скачивали из открытых источников и сами анимировали, поднимали и опускали мостики, создавали разные условия. Получилась своего рода коллаба NeRF и Unreal Engine.

GAN и техника безопасности

Один из компонентов нашей исходной задачи — модель, которая обнаруживает привязь на рабочих. Человек сложнее поезда, состояний у него больше. Поэтому с учётом всех условий мы бы замучились делать тысячи картинок с 3D-человечками в разных позах. Чтобы ускорить процесс, мы использовали GAN — генеративно-состязательную нейросеть.

Модель, которую мы обучаем, находит на изображениях человека и контур привязи.

У привязи разных состояний будет меньше. Поэтому, вместо того, чтобы ставить 3D-человечков в разные позы, мы заставили GAN генерировать изображения людей на основе рандомных контуров привязи. Потом использовали сгенерированные картинки для обучения модели. Эти нейрорабочие выглядели искажёнными, но в целом похожими на людей.

Но сразу возникает вопрос. Чтобы GAN что-то выдавала, ей надо скормить фотографии рабочих с привязями в разных ракурсах — из нашего реального датасета. Не будет ли она копипастить данные из датасета, и в итоге его не расширит? И неужели нельзя обойтись просто аугментацией?

Мы задались этим вопросом и нашли на него ответ:

• F1 Score для модели, при обучении которой мы использовали только аугментацию, равен 82%.

• Аугментация вместе с данными из GAN повысила точность модели до 89%.

А оно того стоит?

Вернёмся к исходной задаче и модели, которая не обучалась на настоящих данных с открытой площадки. Если выпустить её на улицу без дообучения, она будет определять все нужные нам классы с точностью 67%. А после обучения на синтетических данных из игрового движка её точность возросла до 95%. И это модель, которая не видела настоящих фотографий уличной эстакады!

Но нужно ли это вам? Зависит от задачи:

• Для большинства задач хватает открытых источников и настоящих данных с производств. Так, прирост в точности определения привязи на рабочих мы получили за счёт фотографий из гугла.

• Если задача сложная, данных мало и нужно быстрее обучить модель, можно подключить игровой движок. 

• А если модели из игрового движка недостаточно реалистичны, можно использовать NeRF для моделирования реальных объектов с производств.

Значит ли это, что реальные данные не нужны? Нет, мы никогда от них не отказываемся — они всё равно полезны, даже если мы будем их собирать целый год. Постучали по дереву. 

В общем, примерно так мы изворачиваемся ради высокого результата. А вы как решаете проблемы с нехваткой данных? Пишите в комментах, будет интересно обсудить!