Deep Learning — не только котики на мобилках или как мы производим дефектовку тележек локомотивов

hpk85kft0_bi1euqgqvftfm2vic.jpeg

Буквально пару дней назад компания Aurorai передала в опытную эксплуатация систему распознавания дефектов и контроля состояния тележек для локомотивов Ермак. Задача нетривиальная и очень интересная, первым этапом которой было предложено оценить состояние тормозных колодок и ширины бандажа. Нам удалось решить задачу с точность до 1 мм при скорости локоматива до 30 км/ч! Хочу отметить, что благодаря специфики можно было использовать «TTA (test-time augmentation)» — яркий пример kaggle-style хака из соревнований, который плохо ложится на прод и семантическую сегментацию на базе se_resnext50 encoder, которая даёт поразительный по точности результат в предсказании маски.
Описание задачи

Необходимо создать аппаратно-программный комплекс детектирования дефектов тормозных колодок и вывода данных бригадиру смены.

Предпосылки к задаче

Как оказалось, огромное количество колодок, около 80%, меняется в ПТОЛ (пунктах технического осмотра локомотивов), и это происходит каждые 72 часа для каждого локомотива. Основная масса проверок в ПТОЛ — это визуальный осмотр мастером внешней части тележки локомотива.
xplzewvz4vitxelucz2te4liocq.gif

План решения задачи:

  1. Подбор оборудования
  2. Сбор данных
  3. Обучение модели
  4. Разработка сервера с REST API
  5. Разработка клиента для планшета на Android
  6. Проектирование и сборка стойки для размещения камер и света
  7. Опытная эксплуатация

Подбор оборудования

Пожалуй, одной из самых сложных, если не сказать самой сложной, задачей было выбрать камеры, объективы и свет в условиях ограниченного бюджета и времени: MVP нужно было сделать за полтора месяца. За пару дней гугл сделал из меня эксперта по железу для машинного зрения. Выбор остановился на камерах Basler и импульсной подсветке 6к люмен, синхронизирующейся с камерой. В пользу Basler (70 кад/сек, разрешение до 1920×1024) говорило его API на python, что значительно облегчало интеграцию всех компонентов системы, единственный минус — это цена камер ~ 100 т.р.

Выбор объектива к камерам был осложнен непониманием необходимого фокусного расстояния и угла обзора, пришлось рисковать, но вывез калькулятор объективов и щепотка удачи.

Подсветка: опытным путем было установлено необходимое время свечения светодиодов, их тип, параметры линз. Пробовал 3 разные модификации линз для светодиодов, с углом 30, 45, 60 в конечном итоге выбрал матовые линзы с углом 45.
lqvc4a2w5fv7uqqdtveh3a-le3c.jpeg
aptonnrgp3kjkpyyz67jo3g4wnc.jpeg

Сборка и проверка сигнала управления импульсной подсветкой для камеры
n7s_n3chgt_6nfjwv541dluhzn0.jpeg

Для серверного железа я взял, Intel Core i7–7740X Kaby Lake, 46gb RAM, 1 TB SSD и 3×1080Ti — этого вполне хватает, чтобы предсказывать два 3-секционных локомотива не дольше чем за 2 мин.

Колхозное охлаждение бутерброда из видеокарт сдувает 10 град.
f_fbv44tdrie9mdta1djurpv6kw.jpeg

Сбор данных

Создание датасета — это отдельная песня, никому нельзя доверить сие мероприятие, и поэтому меня отправили в далекий малоизвестный городок в глубине нашей необъятной родины. Я сфотографировал на телефон (!!!) около 400 колодок. Забегая вперёд, скажу, что доблестные сотрудники депо, видимо испугавшись ревизора из Москвы, поменяли все колодки на локомотивах на абсолютно новые и покрасили их свеженьким слоем краски, смотреть на это было смешно и страшно. Я предвкушал худшее, правда еще было около 400 фоток совсем других колодок, которые я сделал в московском депо.

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

Ожидание:
buju04qal7c3euollekewwem6ma.jpeg

Реальность:
wjloxc_zlslgjsyag4tubgzg0ms.png
gyxljdwzpmhoxuvbfqie6ixzvpc.jpeg
Тут надо сказать, что не было ни одного примера сильно стертых колодок.

Обучение модели

Лучше всего себя показала модель c энкодером se_resnext50 и декодером с scse блоком из этого репозитория, но scse (реализация для pytorch) пришлось убрать из соображений ускорения процесса предсказания, т.к. нужно было предсказать за минуту. Для обучения модели использовался фреймворк Pytorch 1.0.1, с большим количеством аугментаций от albumentations и самописной аугментаций Horizontal Flip для смены класса при отображение.

def train_transform(p=1):
        return Compose([
            OneOf([
                CLAHE(clip_limit=2),
                IAASharpen(),
                IAAEmboss(),
                RandomBrightnessContrast(brightness_limit=0.8, contrast_limit=0.8),
                HueSaturationValue(hue_shift_limit=50, sat_shift_limit=50, val_shift_limit=50),
                RGBShift(r_shift_limit=50, g_shift_limit=50, b_shift_limit=50),
                JpegCompression(quality_lower=30),
                RandomGamma(),
                GaussNoise()
            ], p=0.3),
            OneOf([
                Blur(),
                MotionBlur(),
                MedianBlur(),
            ], p=0.3),
            ShiftScaleRotate(shift_limit=0.2, scale_limit=0.4, rotate_limit=5, p=0.5),
            Normalize(p=1)
        ], p=p)


В качестве функции потерь я выбрал The Lovász-Softmax loss, вел он себя практически также как и bce + jaccard, но лучше BCE, который слишком сильно фитится на разметку. Так же вызовом стал выбор алгоритма определения порядкового номера колесной пары и колодки, были варианты с metric learning, но мне нужно было быстро показывать результат, и в голову пришла идея разметить колодки на классы 1 и 2, где 1 — ориентация вправо, а 2 — влево. Сеть начала предсказывать не только маску, но ориентацию. С помощью нехитрых эвристик удалось достаточно надежно определять порядковые номера колодок и колёсных пар, далее усредняя предсказания, фактически использую ТТА с небольшим сдвигом объекта при движении и разным углом освещения дает хороший результат в точности маски даже на разрешении 320×320.

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

Пайплайн для обучения взял отсюда MICCAI 2017 Robotic Instrument Segmentation. Процесс обучения состоит из трех этапов: обучения с замороженным энкодером, обучение всей сети и обучение c CosineAnnealingLR. В первых двух стадиях используется ReduceLROnPlateau.

Программирование REST сервера и клиента на Android

Для REST сервера я выбрал flask — проще не придумать, запуск за 2 минуты. Базу данных для хранения решил сделать своими руками в виде простой структуры папок и файла текущего состояния. Приложение для планшета на Android Studio, благо последние версии просто рай для разработчика.

Проектирование и сборка стойки для размещения камер и света

Вспомнил былое время, когда делал зарядные станции для электромобилей, и этот опыт был очень кстати — решили делать из стоек из конструкционного алюминия, напечатанных на 3D-принтере.
syaisdlo96ld_0u_zzrckpnq_ng.jpeg
chpd7sjyqr-x8u99flu2gf_dfy8.jpeg

Приступаем к испытаниям!


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

dkcrt7rtuca4tlosawaygpvw9bw.jpeg
rnzt5amj8om5hiz_w57qxgtxozm.jpeg

Jaccard на валидации доходит до 0.96, визуально колодки от сегментированы очень чётко, добавим усреднение по нескольким фото и получаем очень хорошую точность оценки ширины колодки. По ходу испытаний оказалось, что можно работать и с тележками других локомотивов, но брать более скоростные камеры:
vqzgap76zvv1qlpz9-hqxr8_0s0.jpeg
fuh5cikxtwssvrfy333eyrl8itw.jpeg

В заключении хочу сказать, что технология показала себя очень хорошо и, на мой взгляд, имеет большой потенциал с точки зрения исключения человеческого фактора, снижения времени простоя локомотива и составления прогнозов.

Благодарности

Благодарю сообщество ods.ai, без вашей помощи мне бы не удалось всё это сделать в столь короткие сроки! Огромное спасибо n01z3, надоумевшему меня заняться DL, за его бесценные советы и чрезвычайный профессионализм! Большое спасибо идейному вдохновителю Василию Манько (CEO, компания Aurorai), лучшему дизайнеру Татьяне Брусовой.

До встречи в следующей серии рассказа!

xvx3yvszd_twqjbrtk3rjaxrxpk.png Aurorai, llc

© Habrahabr.ru