Летопись проекта Цельс.Маммография

Всем привет! Полгода назад я рассказывал об опыте написания летописи нашего проекта по маммографии. В комментариях к тому посту вы дружно сказали, что было бы интересно посмотреть на публичную версию этой летописи. Что ж, вот она с некоторыми изменениями - почти без мата (из скриншотов не вырежешь), личных деталей и некоторых архитектурных ноу-хау. Эта летопись охватывает период с осени 2018 года по декабрь 2020.

Глава 1. Начало и баттл

Сентябрь 2018 года. 7 утра в Нью-Йорке, 2 часа дня в Петербурге. Мне звонит Никита Николаев по кличке Качер, мой давний товарищ с универской скамьи.

Начинаем со стандартного обсуждения дел по нашему скоринговому стартапу - что с обновлениями моделек, какие клиенты на горизонте и так далее. В конце, как бы невзначай, Качер бросает - "а как вообще, сложно по рентгену детектировать рак груди?". Мне как раз было интересно глубже погрузиться в тему CV, всё-таки мой основной опыт на тот момент был табличным и NLPшным. В общем, оказалось, что Никитос и так уже предварительно сказал, что мы можем такое сделать, так что моё мнение было не так уж и важно...

Спустя месяц была готова первая версия модели, обученная на датасете из Калужского онкодиспансера с бинарной разметкой - обнаружен у человека рак или нет. Какой-то простой классификатор - VGG16, наверное, сейчас даже и не помню уже. Большая часть месяца ушла на ознакомление с доменной областью и медицинском форматом DICOM. Быстренько упаковал модель во Flask и развернул на Digital Ocean.

Скорость появления прототипа всех порадовала, но вместе с тем быстро стало понятно, что бинарные классификаторы в медицине не взлетят - нужен интерпретируемый детектор или сегментатор. Разметки такой у нас не было, с публичными датасетами тоже напряжёнка, да и разбираться было долго. Так что, в итоге я бессовестно увёл публичные веса архитектуры Faster-RCNN. Одна загвоздочка - они были на первой версии Caffe. Зато пока перегонял их на вторую версию, а затем на PyTorch, очень плотно разобрался в особенностях архитектуры и всех слоях =) Заодно добавил препроцессинг на OpenCV, который удалял лишнюю фигню с груди.

Поверх детектора пришлось накрутить модельку, которая по детектированным объектам определяла общую вероятность наличия злокачественных изменений, на наших данных она выдавала ROC-AUC в районе 0.75. Первая демонстрация врачам прошла успешно, все были в шоке, что какая-то машина способна видеть признаки рака на снимках и даже обводить их коробочкой. Планируем поездку в Калугу для регистрации компании.

Наступил ноябрь, проект официально поехал, а значит пора накинуть немного MLOpsа. Появилась БД, куда сохранялась информация о запросах и предиктах, дашборд на Dash, трекинг экспериментов на MLFlow, алерты на Sentry, сама система была наконец докеризирована. Заодно запиливаю на Flask простенькое веб-приложение для демонстрации системы в лайве.

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

В мае я провёл первые собесы (почти все ребята из первой волны найма до сих пор с нами, люблю вас). В июне я собрал свои пожитки и вернулся из Нью-Йорка в Петербург открывать офис в подвальном помещении на Обводном канале. В репе появляются первые коммиты не от crazyfrogspb, заводим свой Слак и продолжаем дообучать модельку. Перекатываемся на реализацию Faster-RCNN от torchvision, а я нечеловческими темпами потребляю литературу по Agile и тимлидству.

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

В команде проекта появляются не только ML-специалисты - например, фронты наконец-то нарисовали нормальный интерфейс вместо моего убогого фласкового.

В августе продолжаем релизить улучшенные версии модели, а параллельно совершенствовуем архитектуру проекта. Монолит распиливается на CPU и GPU-части по этому мануалу, пишем первые автоматизированные метрик-тесты, улучшаем препроцессинг изображений.

Осенью начинаем активную разметку на платформе Supervisely, а значит вскоре система научится детектить новые классы - ФКМ, утолщения кожи, соски. Соски начали размечать, потому что система часто воспринимала их как образования, такой вот hard negative mining. Начинаем писать документацию, например, появляется такая вот схемка:

Осенью активно экспериментируем с архитектурой - пробуем STN, заменяем батчнорм на групнорм из-за маленького батч-сайза. Команда разметки вместе с ML запиливает первый рулбук по разметке для врачей.

Тем временем, в России начинает активно развиваться индустрия ИИ в рентгенологии. Основная активность идёт, конечно же, из Москвы - Сергей Павлович Морозов и его команда анонсируют ИИ-баттл, который пройдёт в рамках традиционной конференции рентгенологов. Этот баттл станет предтечей знаменитого Московского эксперимента.

В ноябре и декабре идёт активная разработка новой архитектуры, которая призвана заменить жирную систему, которая на тот момент включает аж три раздельных детектора. Главная проблема на тот момент - в разных датасетах размечены разные классы, эта проблематика впоследствии будет хорошо описана в этой статье. Но на тот момент она ещё не была написана, так что пришлось придумывать всё самим =)

В день перед баттлом получаем датасет, который надо обработать (100 исследований). Идут жаркие споры по поводу того, стоит ли использовать новую, не вполне оттестированную модель. Я пытаюсь всех убедить, что старая модель будет работать плохо - например, в ней нет нового препроцессинга, который корректно работает на разнообразных исследованиях с разного оборудования. Эмоциональное убеждение работает, так что мы рискуем и за минуту до дедлайна отправляем архив с предсказаниями и визуализациями.

На баттле всё проходит хорошо, и мы забираем гран-при. Это был первый крупный успех в истории Цельса, и вся ML-команда в полном составе выезжает в Калугу на лютое празднование.

Глава 2. Московский эксперимент

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

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

К этому моменту у руководства компании появляется чёткое понимание, что пора бы уже не только тратить, но и начать зарабатывать деньги. Наши наивные надежды быстро начать рубить капусту в регионах с очень сырым продуктом не оправдались. На наше счастье появился он - Московский эксперимент. Процесс интеграции в Эксперимент с самого начала был весьма непростым - анкета, очное собеседование в Москве, техническая интеграция с ЕМИАС, функциональное и наконец калибровочное тестирование с измерением метрик точности.

Дата калибровки всё откладывалась из-за техническим проблем, но вот наконец этот момент настал. Всё прошло нормально, все исследования обработались, и 26 мая 2020 года я с замиранием сердца щёлкнул на PDF с результатами. ROC-AUC - 0.73, при минимальном требовании в 0.81. Настоящая пропасть в 8 пунктов. Состоялся очень тяжёлый разговор с топ-менеджментом, до сих пор помню тягостное ощущение после того созвона. Я сидел на диване в переговорке в новом офисе на Невском, и паника начинала захватывать мои мысли. Совместным решением был установлен очень жёсткий для DL-проектов срок на исправление ситуации - две недели. В случае неуспеха в теории могла встать речь и о закрытии проекта. Началась так называемая Миссия Сапсан.

Через две недели бессонных ночей и моментов полного бессилия и апатии рождается новая версия системы. Её главное изменение связано с добавлением второго канала с хитрым препроцессингом в качестве дополнительного инпута. Вскоре после этого мы узнали о так называемой VOI-LUT трансформации. Оказывается, внутри дайкома зашита информация о том, как нужно преобразовать значения пикселей, чтобы отсечь пиксели, которые неинформативны для данной модальности. Это позволило нормально обрабатывать так называемые “снимки с засветами”, которые успели изрядно надоесть к этому моменту. Препроцессинг в целом продолжает обрастать хаками и эвристиками, а мы решили ещё раз детально изучить особеннности DICOM-формата.

Параллельно идёт работа над улучшением риск-модели, которая присваивает общую вероятность рака всему исследованию, ведь главный показатель калибровки - это как раз ROC-AUC бинарной классификации. Докидываем фичей, калибруем вероятности (внезапно узнали, что их по правилам на бэке округляют до второго знака).

Молимся. Релизимся. 26 июня получаем результаты повторной калибровки. Ровно 0.81. Это ПОБЕДА!!!

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

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

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

Этот год навсегда останется в истории Цельса и проекта ММГ - это год, в который мы упали, но поднялись - и взлетели!!! Но о взлёте в другой раз 😊