Буквально пару дней назад компания Aurorai передала в опытную эксплуатация систему распознавания дефектов и контроля состояния тележек для локомотивов Ермак. Задача нетривиальная и очень интересная, первым этапом которой было предложено оценить состояние тормозных колодок и ширины бандажа. Нам удалось решить задачу с точность до 1мм при скорости локоматива до 30 км/ч! Хочу отметить, что благодаря специфики можно было использовать “TTA (test-time augmentation)” – яркий пример kaggle-style хака из соревнований, который плохо ложится на прод и семантическую сегментацию на базе se_resnext50 encoder, которая даёт поразительный по точности результат в предсказании маски.
Описание задачи
Необходимо создать аппаратно-программный комплекс детектирования дефектов тормозных колодок и вывода данных бригадиру смены.
Предпосылки к задаче
Как оказалось, огромное количество колодок, около 80%, меняется в ПТОЛ (пунктах технического осмотра локомотивов), и это происходит каждые 72 часа для каждого локомотива. Основная масса проверок в ПТОЛ – это визуальный осмотр мастером внешней части тележки локомотива.
План решения задачи:
- Подбор оборудования
- Сбор данных
- Обучение модели
- Разработка сервера с REST API
- Разработка клиента для планшета на Android
- Проектирование и сборка стойки для размещения камер и света
- Опытная эксплуатация
Подбор оборудования
Пожалуй, одной из самых сложных, если не сказать самой сложной, задачей было выбрать камеры, объективы и свет в условиях ограниченного бюджета и времени: MVP нужно было сделать за полтора месяца. За пару дней гугл сделал из меня эксперта по железу для машинного зрения. Выбор остановился на камерах Basler и импульсной подсветке 6к люмен, синхронизирующейся с камерой. В пользу Basler(70 кад/сек, разрешение до 1920х1024) говорило его API на python, что значительно облегчало интеграцию всех компонентов системы, единственный минус — это цена камер ~ 100 т.р.
Выбор объектива к камерам был осложнен непониманием необходимого фокусного расстояния и угла обзора, пришлось рисковать, но вывез калькулятор объективов и щепотка удачи.
Подсветка: опытным путем было установлено необходимое время свечения светодиодов, их тип, параметры линз. Пробовал 3 разные модификации линз для светодиодов, с углом 30, 45, 60 в конечном итоге выбрал матовые линзы с углом 45.
Сборка и проверка сигнала управления импульсной подсветкой для камеры
Для серверного железа я взял, Intel Core i7-7740X Kaby Lake, 46gb RAM, 1 TB SSD и 3х1080Ti — этого вполне хватает, чтобы предсказывать два 3-секционных локомотива не дольше чем за 2 мин.
Колхозное охлаждение бутерброда из видеокарт сдувает 10 град.
Сбор данных
Создание датасета – это отдельная песня, никому нельзя доверить сие мероприятие, и поэтому меня отправили в далекий малоизвестный городок в глубине нашей необъятной родины. Я сфотографировал на телефон(!!!) около 400 колодок. Забегая вперёд, скажу, что доблестные сотрудники депо, видимо испугавшись ревизора из Москвы, поменяли все колодки на локомотивах на абсолютно новые и покрасили их свеженьким слоем краски, смотреть на это было смешно и страшно. Я предвкушал худшее, правда еще было около 400 фоток совсем других колодок, которые я сделал в московском депо.
Оставалось только верить в чудо, навалить аугментаций, придумать эвристик для удаления ошибочных сегментов, коих было не мало, так как я не подумал об анти примерах.
Ожидание:
Реальность:
Тут надо сказать, что не было ни одного примера сильно стертых колодок.
Обучение модели
Лучше всего себя показала модель 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-принтере.
Приступаем к испытаниям!
Результат превзошел все ожидания. Специалистам компьютерного зрения задача может показаться достаточно прямолинейной и простой. Однако у меня был некоторый скепсис из-за двух вещей: во-первых, обучающая выборка была небольшой и не содержала граничные кейсы типа очень тонких колодок; во-вторых, испытания проходили в сильно других условиях съемки и освещения.
Jaccard на валидации доходит до 0.96, визуально колодки от сегментированы очень чётко, добавим усреднение по нескольким фото и получаем очень хорошую точность оценки ширины колодки. По ходу испытаний оказалось, что можно работать и с тележками других локомотивов, но брать более скоростные камеры:
В заключении хочу сказать, что технология показала себя очень хорошо и, на мой взгляд, имеет большой потенциал с точки зрения исключения человеческого фактора, снижения времени простоя локомотива и составления прогнозов.
Благодарности
Благодарю сообщество ods.ai, без вашей помощи мне бы не удалось всё это сделать в столь короткие сроки! Огромное спасибо n01z3, надоумевшему меня заняться DL, за его бесценные советы и чрезвычайный профессионализм! Большое спасибо идейному вдохновителю Василию Манько (CEO, компания Aurorai), лучшему дизайнеру Татьяне Брусовой.
До встречи в следующей серии рассказа!
Aurorai, llc