Всем привет! Мы — научно‑исследовательская лаборатория «Сильный ИИ в медицине» в Институте AIRI. Наша группа разрабатывает решения на основе искусственного интеллекта в медицине. На днях стало известно, что мы заняли призовое, пятое место в конкурсе «MIDRC XAI Challenge: Decoding AI Decisions for Pneumonia on Chest Radiographs» с опытом участия в котором, мы хотели бы поделиться.
Пару слов о конкурсе
Несколько лет назад COVID-19 изрядно потрепал нервы всем, особенно врачам. Рук в больницах в то время не хватало. Помощь в виде автоматизации хотя бы части забот была бы очень кстати — например, диагностики пневмонии, ассоциированной с коронавирусной инфекцией, по снимкам рентгена лёгких. С этой целью в США был создан Центр медицинских изображений и данных MIDRC.
В постковидную эпоху его задачи были расширены на содействие исследованиям, в которых машинное обучение используется для обнаружения, диагностики, мониторинга и прогнозирования как можно большего числа заболеваний, не только COVID-19. Чтобы привлечь в эту область ML‑комьюнити, в MIDRC организовали соревнование MIDRC XAI Challenge: Decoding AI Decisions for Pneumonia on Chest Radiographs.
Этот конкурс был направлен на решение одной из ключевых задач в области анализа медицинских изображений: создание интерпретируемых и надежных моделей искусственного интеллекта. Задачей участников была разработка и обучение модели классификации фронтальных рентгенограмм грудной клетки на наличие затемнений в лёгких (lung opacity), связанных с любым типом пневмонии. Результатом работы модели должны были стать:
вероятность того, что пациент страдает пневмонией любого типа,
карта «объяснимости», интерпретируемая как вероятность наличия затемнений в лёгких в каждом пикселе (того же размера, что и входное изображение).
Организаторы не предоставляли датасет для обучения. В качестве метрики, по которой проводилась оценка качества модели, организаторы выбрали weighted log loss. При этом вес пикселя определялся следующим образом: у организаторов была команда из 18 врачей‑экспертов, каждое исследование из валидационной выборки было размечено тремя экспертами. Они рисовали сегментационную маску зон затемнений, далее эти маски накладывались друг на друга. Пиксели сегментационной маски, размеченные тремя экспертами, имели вес 45, двумя — 30 и при разметке только одним экспертом — 15.
Посмотрев на такой сетап, мы принялись за работу.
Метрики
Первое, на что мы обратили внимание, это выбор метрики организаторами. Известно, что log‑loss — плохая метрика, чувствительная к ошибкам в области предельных значений, а значит по её значению трудно интерпретировать качество моделей. Действительно, если модель предсказывает в каком‑нибудь пикселе 0, а там на самом деле 1, то лосс взорвется, поэтому мы делали сглаживание этих предельных значений.
Оказалось, что без сглаживания значений вероятности мы получали log‑loss хуже, чем случайная модель (а она, по словам организаторов, давала log‑loss = 1)! И это несмотря на то, что модель рисовала разумные сегментационные маски и выдавала ROC‑AUC = 0.9.
В общем, мы поняли, что для работы над моделями нам нужно что‑то другое в роли оценки. Чтобы сравнивать эксперименты между собой, не привязываясь к постобработке, мы выбрали следующие сегментационные метрики, основанные на Intersection over Union (IoU):
IoU по маскам, где есть разметка минимум 1 эксперта (gt_mask = sum(masks) >= 1),
IoU по маскам, где есть разметка минимум 2 экспертов (gt_mask = sum(masks) >= 2),
IoU по маскам, где есть разметка всех 3 экспертов (gt_mask = sum(masks) >= 3).
Такой набор метрик для трёх независимых разметок от экспертов помогал лучше понимать, насколько консервативна обученная модель, и выдвигать гипотезы для следующих экспериментов.
Забегая вперед, отметим, что мы, конечно же, рассчитывали метрику соревнования (log loss) для сегментационных моделей, но лишь на этапе сбора финального ансамбля — для расчета коэффициентов, с которыми стоит смешивать предсказания.
Данные
Как уже было отмечено выше, датасета для обучения не было, однако организаторы предоставили 14 валидационных рентгенов грудной клетки, которые можно было использовать в качестве «калибровочных» (посмотреть на примеры разметки, потестить код инференса и код расчета метрик). Среди 14 размеченных врачами рентгенов три были без затемнений, а одиннадцать с затемнениями, к которым шла разметка экспертов (три разметки на один снимок). Все рентгены были предоставлены в формате DICOM (Digital Imaging and Communications in Medicine) — это стандарт обработки, хранения, передачи, печати и визуализации медицинских изображений.
На начальном этапе мы проанализировали предоставленные изображения, чтобы понять их структуру и особенности. Для точного выделения областей затемнений на изображения накладывались маски, что позволяло изолировать эти участки и исключать лишние детали, которые могли бы внести шум в дальнейший анализ. В итоге данный этап подготовил изображения для обсуждения их с врачами, что способствовало более глубокому пониманию поставленной организаторами задачи.
В ходе обсуждения с врачами‑рентгенологами были получены следующие комментарии:
Существует несколько глоссариев терминов для торакальной визуализации, куда включены определения важных для визуализации симптомов, анатомических структур и наиболее частых патологических состояний, и они отличаются друг от друга.
По глоссарию Fleischner Society: Glossary of Terms for Thoracic Imaging к затемнениям в лёгких можно также относить такие заболевания как пневмония и ковид.
Имеет место низкая согласованность масок затемнений между врачами, связанная с тем, что четкие границы затемнений выделить крайне сложно. В этом случае врачи на свое усмотрение могут выделить как лёгкое целиком, так и только очаги затемнений. Некоторые маски на валидационных кейсах вызвали споры у наших врачей‑экспертов.После встречи мы сделали следующие выводы:
После встречи мы сделали следующие выводы:
Из целого класса причин, которые способны вызвать затемнение (рак, ателектазы, диссеменации и много ещё чего), нам стоит сфокусироваться на пневмонии и ковиде.
Четко определить границы затемнений крайне тяжело и не всегда возможно.
Некоторые валидационные кейсы имеют спорную разметку.
После общения с экспертами мы выполнили анализ данных (Exploratory Data Analysis, EDA), который помог изучить их основные закономерности. Мы проанализировали согласованность врачей в рамках валидационных кейсов и получили такие результаты:
Он говорит нам, что среди представленных рентгенов есть несколько кейсов, с высокой согласованностью экспертов (все 3 врача выбрали одну и ту же область). Но в общем случае врачи не могут прийти к единому мнению для определения области lung opacity. Эти факты привели нас к выводу, что итоговую область затемнений стоит размывать для улучшения метрик качества.
Дальше мы решили проанализировать DICOM‑тэги изображений для поиска ликов и дополнительной информации, которая могла бы оказаться полезной. Важная особенность DICOM — это объединение визуальных данных (изображений) с мета‑информацией, хранящейся в тегах.
Примером полезной информации, которая может быть использована, является описание модальности снимков, а также мета‑информация о проводимом исследовании. Наш поиск показал, что часть исследований косвенно относятся к датасету RSNA Covid-19 Dataset Default. Мы решили использовать для обучения датасеты: VinDr‑CXR, RSNA Pneumonia, QaTa‑COV19, Mendeley Tomsk, COVID-19 Radiography Dataset, SIIM‑FISABIO‑RSNA‑COVID-19.
Отметим, что к конкурсу мы подключились в конце этапа валидации за месяц до конца соревнования, и, поскольку решений у нас к этому моменту не было (не считая сырой сегментационной модели), мы остались только с 14 размеченными кейсами, без возможности провалидировать наши метрики где‑то еще, не считая собранных датасетов.
Идеи
Посидев и хорошо подумав, мы набросали несколько идей. Решили, что наше итоговое решение будет состоять из следующих этапов:
Бинарная классификация на наличие lung opacity.
Сегментация затемнений в лёгких на разрешении 336×336.
Прогон чувствительной сегментационной сеткой внутри баундинг боксов с предыдущего шага.
Сегментация лёгких.
Далее по каждому из пунктов поподробнее.
1. Бинарная классификация
Задача, казалось бы, простая — взять датасеты и обучить модель, что мы и сделали. Обучили EfficientNetB0 на COVID-19 Radiography Dataset. Модель хорошо различала норму и ненорму и при этом также очень хорошо разделяла затемнения, пневмонию и ковид. Параллельно мы искали опенсорсную модель, чтобы сравниться с нашей кастомной.
Нашли и затестили модель из библиотеки torchxrayvision. Несмотря на то, что AUC по классу «lung opacity» у нее был 0.92 по сравнению с 0.96 на нашей кастомной модели, мы решили оставить опенсорсную модель. Интуиция была в том, что она обучалась на существенно бо́льшем наборе данных (RSNA Pneumonia Challenge, NIH chest X‑ray8, PadChest (University of Alicante), CheXpert (Stanford), MIMIC‑CXR (MIT)) и теоретически должна была выдавать стабильный результат на любых данных. Все 14 валидационных кейсов от организаторов она верно классифицировала.
2. Сегментация затемнений в лёгких
Вариантов, как построить сегментационную модель, у нас было несколько.
FPN from CLIP-ConvNext
Первый вариант — это модель на основе нашего секретного оружия — in‑house CLIP. Архитектурно это FPN, где encoder инициализируется из convnext, предобученного методом CLIP на большом количестве парных данных с рентгенами и их описаниями.
Понятно, как использовать датасеты с масками. Но что делать с боксами? Самое ленивое решение — закрасить их. Интуиция в пользу такого ленивого подхода опиралась на вид разметки калибровочного датасета: у какого‑то эксперта закрашенная область больше, у какого‑то меньше, а само понимание точной границы патологии субъективно и не согласовано между радиологами. То есть мы подсовываем в обучение разные типы разметки (по сути, разные мнения радиологов), в надежде, что модель сама разберется и на краях немного сгладит грубые разметки прямоугольником. Так и произошло.
Сначала мы взяли для экспериментов все доступные датасеты, в том числе детекционные VinDR и RSNA Pneumonia. Увы, при этом происходил какой‑то domain shift / bias (называйте как хотите) — при обучении на датасете из одного источника (например Mendeley + QataCov19) метрики были хорошие, а на отложенном тесте (например) VinDR сильно проседали.
Поэтому нам стало понятно, что датасеты нужно миксовать (иначе можно долго копаться в причинах шифтов и баесов). В микс вошли RNSA‑Pneumonia, QataCov19, Mendeley. А от Vindr при обучении этой модели мы в итоге отказались из‑за недоверия к протоколам разметки (которые мы обсуждали с врачами выше) и просадки по метрикам. Такой выбор датасетов давал стабильные метрики по разным подвыборкам, а также разумные значения тестового ROC AUC. C псевдолейбелингом для данной модели решили не заморачиваться, так как мы планировали, что эта работающая в невысоком разрешении модель станет лишь предварительным этапом для следующей, более чувствительной и точной сегментационной сетки.
В результате для данного варианта модели мы тупо закрашивали боксы в датасетах с разметкой боксов (RSNA Pneumonia) и замешивали их с сегментационными примерами. К нашему удивлению, на предварительной валидации на калибровочном датасете получившаяся модель давала самый лучший лосс. Она была менее консервативной и как бы включала в себя каждого отдельного эксперта (которые могли между собой не пересекаться), в то время как описанные далее nnUnet‑ы были более консервативны — более уверенно обводили области, где нашли затемнение хотя бы два эксперта. Отметим, что финальная модель отличалась от промежуточных экспериментов более агрессивной аугментацией картинок (в разумных пределах гиперпараметров).
nnU-Net-v2 + box pseudolabeling
Второй вариант — это использование фреймворка nnU‑Net (v2) в режиме 2D. Данная модель была обучена в два этапа.
Сначала было обучение сегментацинной модели (также nnU‑Net) на кропах патологий из датасетов с масками (Mendeley Tomsk, QaTa‑COV19) и доразметка (aka pseudo‑labeling) с ее помощью датасетов, содержащих только боксы (VinDr‑CXR, SIIM‑FISABIO‑RSNA‑COVID-19), то есть для каждого бокса генерировалась маска. В результате этого этапа мы получили все данные с масками.
Дальше мы проводили обучение сегментационной модели, которая обрабатывает все изображение целиком. В процессе обучения были использованы разные комбинации датасетов и размеров изображений, каждая из которых оценивалась на валидационном наборе из 14 снимков. Лучшей конфигурацией nnUnet оказался следующий набор: Mendeley Tomsk + QaTa‑COV19 + VinDr‑CXR, бо́льшая сторона снимков приводилась к 768 пикселям с сохранением соотношения сторон.
Впоследствии сегментационная модель, обученная для псевдо‑разметки, была также использована на следующем этапе для уточнения масок внутри боксов.
nnU-Net-v2 + experiments on preprocessing strategies
Третий вариант — использовать nnUnet, варьируя лишь методологию работы с данными, например, препроцессинг, аугментацию или выбор датасетов. Для обучения модели сегментации пневмонии были использованы два публично доступных сегментационных датасета: Mendeley Tomsk и QaTa‑COV19. Валидация проводилась на всё тех же 14 снимках, но здесь уже мы смотрели на основную метрику соревнования — взвешенный log loss.
Изначально, чтобы сбалансировать выборку, добавив в нее кейсы без пневмонии, мы попытались использовать дополнительные датасеты, VinDr‑CXR или SIIM‑FISABIO‑RSNA‑COVID-19. По доступной в них классификационной или детекционной разметке можно отфильтровать кейсы без пневмонии. Но по валидационным результатам качество сегментации со сбалансированной и расширенной выборкой было хуже, поэтому от их использования при подготовке финальной модели мы отказались.
Поэкспериментировали мы и с размерами. Перед обучением все снимки были переведены в grayscale RGB и отмасштабированы к размерам изображения N ✕ N. N мы выбрали из вариантов {336, 512, 768} по наилучшим валидационным результатам и зафиксировали N=336.
Мы также проверили гипотезу о применении эквализации. Эквализация или, более полно, приведение к эквализированному распределению интенсивностей (histogram equalization) — это один из популярных методов препроцессинга медицинских снимков, CXR в частности. Но сравнив по валидационным результатам модель с эквализацией и без неё, в том числе и на разных разрешениях изображений, мы приняли решение от неё отказаться.
Нельзя, однако, не признать, что эквализация действительно делает некоторые снимки визуально отчетливее (для человека), но может портить результат в среднем. Поэтому мы решили проверить идею о добавлении эквализации в качестве случайной аугментации снимков с вероятностью p=0.5 на этапе обучения. Это дало определенный прирост по валидационным метрикам, поэтому в обучение финальной модели мы включили этот метод аугментации.
Наконец, мы попытались применить маску лёгких, полученную с помощью torchxrayvision для препроцессинга. Выглядело это примерно так:
К сожалению, этот трюк значительно негативно сказался на валидационных метриках. Возможно, так вышло из‑за специфики разметки пневмонии — многие контуры, размеченные экспертами (в том числе и с соревнования), выходят за пределы маски лёгких.
3. Прогон чувствительной сегментационной сеткой внутри баундинг боксов
Как уже упоминалось ранее, мы решили попробовать в качестве дополнительной сегментационной сетки использовать модель, работающую с кропами. По нашей задумке, она, работая в более высоком разрешении, должна улучшать предсказания более грубых моделей с предыдущего шага.
На практике выяснилось, что обученная с помощью nnUnet чувствительная сетка показала низкие метрики во время прогона всего пайплайна. Мы связываем это изменением того, как выбираются боксы. Возможно, стоило повторить процедуру с моделью на основе нашего CLIP.
На инференсе эта модель всё же использовалась, но только если по результатам предыдущих шагов было выявлено наличие пневмонии.
4. Сегментация лёгких
Сегментация лёгких нужна нам потому, что логично искать местоположение очагов затемнения именно там. Её, однако, постигла та же участь, что и классификацию.
Сначала мы учили nnUnet и MedSam на датасете COVID-19 Radiography Dataset. На нем все получалось идеально, а вот на валидационных кейсах от организаторов модель сыпалась. На помощь пришел все тот же torchxrayvision, который дал хороший dice на самом датасете и визуально отличную сегментацию на кейсах от организаторов.
Решение
В результате итоговое решение выглядело так:
Шаг № 1. Классификация на наличие затемнений в лёгких с помощью общедоступной модели, взятой из пакета torchxrayvision.
Шаг № 2. Сочетание 3-х моделей сегментации. Они, как уже было описано выше, обучались на общедоступных данных: QataCov19, наборе данных по Covid от Mendeley, VinDr и RSNA Pneumonia.
В итоге после первых двух шагов мы имеем 4 модели классификации наличия патологии. Их среднее показывает, есть ли признаки затемнений лёгких на изображении или нет. Если нет, мы переходим к последующей обработке снимков лёгких; если да, мы переходим к шагам 3, 4 и 5.
Шаг № 3. Объединение масок из трех моделей сегментации. Связанные компоненты из этих масок извлекаются и обрабатываются другой моделью сегментации, тоже nnUNet, которая была обучена на кропах, содержащих затемнения в лёгких. Этот шаг уточняет предыдущие маски, так как они были получены с относительно низким разрешением.
Шаг № 4. Поиск лёгких на изображении сеткой из torchvision.
Шаг № 5. Смешиваем четыре результата сегментации и проводим последующую обработку по отдельности. Мы рассматриваем каждую модель сегментации как экспертную. Существует несколько подходов к смешиванию промежуточных результатов.
Подход, который мы условно обозвали «Зоны» работает так: чтобы сгладить значения на выходной карте вероятностей, мы присваиваем разные значения областям, отмеченным четырьмя последовательными прогонами четырех моделей. Значения в этих областях оптимизируются во время внутренней проверки. «Зоны» сильно вдохновлены разметкой, предоставленной организаторами, в которой были достаточно большие области, размеченные только одним, только двумя или только тремя врачами.
Подход под названием «Взвешенный» устроен иначе: мы берем карты вероятностей и суммируем их вместе с весами, которые обратно пропорциональны их независимым взвешенным логарифмическим значениям потерь (рассчитанным на основе 14 исследований на этапе калибровки). Это дает нам гладкую карту вероятностей.
После этого, чтобы еще больше сгладить небольшие вероятности, мы применяем постобработку с использованием расширенной маски лёгких, увеличивая значения пикселей внутри этой маски. Мы выполняем это дважды с разными значениями пикселей лёгких.
Всего у нас было три попытки, которые мы использовали так:
Попытка 1: «Зоны» с более высокими значениями пикселей лёгких (параметр np.clip, равный 0.3 для положительных кейсов, 0.1 — для отрицательных кейсов, и фон, всегда равный 0.05)
Попытка 2: «Взвешенная» сумма вероятностей + постобработка лёгких, как в попытке 1.
Попытка 3: «Зоны» с более низкими значениями пикселей лёгких (0.15 для положительных кейсов, 0.05 для отрицательных кейсов, фон равен 0.05 для положительных и 0 для отрицательных кейсов).
Какая‑то из этих попыток (какая именно, нам пока не сообщили) сработала и принесла нам 5 место на соревновании! Было бы также интересно узнать, из чего состоят пайплайны победителей, но на момент написания этого текста обещанный организаторами общий звонок так и не состоялся.
Заключение
Участие в MIDRC XAI Challenge позволило нам глубже понять задачи автоматизации медицинской диагностики. Мы извлекли ценные уроки о значении согласованности экспертных данных, сложности интерпретации результатов и важности тщательной проработки всех этапов — от выбора архитектуры до постобработки.
Вместе с этим наша команда прошла и проверку на стрессоустойчивость. Утомительный сбор датасетов и проверка гипотез, кривой портал MIDRC и горящие дедлайны здорово нас закалили!
Надеемся, наш опыт будет полезен читателям, которые дочитали этот пост до конца. Если у вас остались вопросы, можете задать их в комментариях или написать на почту Bespalov@airi.net.
Желаем всем удачи и отличных метрик!