Неоднозначные выводы о ROI в УЗИ классификации

В медицинском компьютерном зрении есть идея, перед которой трудно устоять: сначала найти патологический объект, а потом классифицировать уже не весь снимок, а только его. Для УЗИ это звучит почти как здравый смысл. В полном кадре хватает всего, что модели, казалось бы, видеть не нужно: подписи аппарата, измерительные маркеры, шум, лишний фон. Логика простая: берём маску опухоли, оставляем область интереса, всё остальное закрашиваем — и даём классификатору «чистую» картинку.

Но с УЗИ быстро выясняется, что «лишнее» — не такая очевидная категория. Клиницист смотрит не только на то, что находится внутри образования. Ему важны контур, стенка, перегородки, акустические эффекты, соседние ткани, масштаб и даже положение объекта в кадре. Поэтому жёсткая обрезка по маске может сработать не как фильтр от шума, а как ножницы: вместе с фоном мы рискуем отрезать часть диагностического сигнала.

Мы проверили это на открытом наборе данных MMOTU / OTU-2D с УЗИ-изображениями опухолей яичника. Суммарно мы обучили 64 модели с различными вариантами обучающих данных, обучение проводили в течении 8 часов на видеокарте A100 в Google Colab.

Коротко:

• всего получено 72 результата: 64 обучения отдельных моделей и 8 ансамблей;

• в задаче на 7 классов полный кадр дал macro-F1 = 0.4971 ± 0.0185, а вырезка по маске с занулением фона дала 0.6052 ± 0.0231;

• в задаче на 2 класса полный кадр дал macro-F1 = 0.6752 ± 0.0266, а вырезка по маске с занулением фона дала 0.7236 ± 0.0222;

• ансамбль восьми вариантов входа оказался сильнее любого отдельного входа: 0.7025 ± 0.0123 для 7 классов и 0.7728 ± 0.0303 для 2 классов;

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

Одна и та же разметка может породить разные входы для классификатора. Это принципиально разные гипотезы о том, где находится полезная для классификации информация.

Что именно проверяем

Важно сразу сузить постановку. Модель не решает скрининг, не ищет яичник в потоке кадров, не оценивает качество исследования и не учитывает анамнез. Мы берем уже подготовленные изображения из MMOTU / OTU-2D: каждый пример относится к области опухолей или образований яичника и имеет метку класса. Модель работает только с изображениями, где точно есть образование яичника.

Формулировка эксперимента такая:
— есть УЗИ-изображение образования яичника; нужно определить тип образования по изображению.

ROI это region of interest, т.е. область интереса вокруг размеченного образования. Мы используем этот термин в широком смысле: не только «пиксели внутри маски», но и разные способы показать модели объект и ближайшее окружение. Это важное различие. ROI можно сделать мягким, оставив прямоугольник вокруг объекта, а можно сделать жестким, занулив все вне маски. Результаты у этих вариантов разные и можно их посмотреть на фотке ниже.

Данные и классы

Основной набор данных в эксперименте — MMOTU / OTU-2D. Это публичный набор УЗИ-изображений опухолей яичника, опубликованный на Figshare.

В статье MMOTU описаны 1469 двумерных УЗИ-изображений (после нашей чистки мы оставили только 1202 изображения) и 170 изображений контраст-усиленного УЗИ с разметкой на уровне всего изображения и на уровне пикселей.

Разбиение такое:

Для двухклассовой постановки эти же изображения просто сворачиваются в две группы:

Это группировка для проверки влияния входного представления. Она нужна, чтобы проверить: одинаково ли ROI влияет на тонкую 7-классовую задачу и на более грубую задачу — простая киста против всего остального.

2-классовая задача выглядит проще, но она резко несбалансирована: в тесте только 19 простых кист против 363 остальных случаев. Поэтому обычная доля правильных ответов здесь может выглядеть высокой даже при слабой работе по редкому классу. Основная метрика в статье — macro-F1.

Какие входы видел классификатор

Из одного и того же изображения и одной и той же маски были собраны восемь вариантов входа:

В таблицах ниже есть ещё строка «ансамбль восьми входов». Это не девятый вариант входа. Это усреднение вероятностей восьми моделей, обученных на восьми разных входах при одном и том же повторе.

Дизайн эксперимента

Чтобы сравнение было честным, архитектура везде одна:

*Пунктиром показана функция потерь на обучении, сплошной линией — на валидации.

Результаты

Столбцы показывают средний macro-F1 по четырём повторам, чёрные отрезки — стандартное отклонение.

Каждая точка — один повтор. Чёрная риска — среднее по четырём повторам. Эти графики показывают разброс, который скрывается за средними значениями.

7 классов

В задаче на 7 классов результат однозначный: все варианты, использующие маску или вырезку вокруг неё, оказались лучше полного кадра во всех четырёх повторах. Лучший одиночный вход — вырезка по маске с занулённым фоном: 0.6052 ± 0.0231 против 0.4971 ± 0.0185 у полного кадра.

Прямоугольные вырезки тоже помогают, но слабее точной маски. При этом самый большой отступ 80% даёт меньший выигрыш: объект снова становится частью крупного фрагмента кадра, и модель получает больше посторонней информации.

Полный кадр с маской как дополнительным каналом улучшил результат относительно обычного полного кадра, но не обошёл вырезку по маске.

Ансамбль восьми входов дал лучший общий результат: 0.7025 ± 0.0123. Он сильнее любого одиночного варианта, потому что разные входы сохраняют разные признаки и ошибаются по-разному.

По дополнительным метрикам картина согласуется с macro-F1: у полного кадра balanced accuracy = 0.4933, у вырезки по маске 0.6303, у ансамбля 0.6947. Для macro-AUROC соответственно 0.8338, 0.8883 и 0.9226.

2 класса

В 2-классовой задаче картина менее жёсткая. Полный кадр не провалился, а многие прямоугольные вырезки отличаются от него всего на несколько сотых macro-F1. Часть повторов у прямоугольных вариантов хуже полного кадра, поэтому здесь нельзя честно сказать: «любой ROI-вход точно лучше».

Но два вывода устойчивы.

Первый: вырезка по маске с занулённым фоном снова стала лучшим одиночным входом: 0.7236 ± 0.0222.

Второй: ансамбль восьми входов дал 0.7728 ± 0.0303 и был лучше полного кадра во всех четырёх повторах.

По AUROC полный кадр дал 0.8298, вырезка по маске — 0.8733, ансамбль — 0.9175. У прямоугольника с отступом 50% AUROC = 0.9043, но macro-F1 у него нестабилен: стандартное отклонение 0.0681. Поэтому я не считаю его лучшим одиночным вариантом.

Что означает результат

Для 7-классовой задачи полный кадр оказался слишком шумным входом для небольшой EfficientNet-B0. Модель лучше работала, когда ей показывали образование крупнее и убирали фон. Особенно это заметно на редких классах.

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

Полный кадр плюс маска оказался промежуточным вариантом. Он лучше полного кадра, но хуже жёсткой вырезки по маске. Значит, дополнительный канал с маской не гарантирует, что модель будет использовать разметку самым полезным образом. Иногда лучше заранее изменить само изображение.

Ансамбль сильнее, потому что разные входы сохраняют разные части информации (далее наши гипотезы, которые ещё только предстоит проверить):

• полный кадр сохраняет масштаб и общую сцену;

• вырезка по маске фокусируется на самом образовании;

• прямоугольники с разным отступом оставляют разные уровни окружения;

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

Если отдельный вход теряет часть информации, другой вход может её сохранить. Поэтому ансамбль восьми входов в 7-классовой задаче поднимает macro-F1 до 0.7025, а в 2-классовой — до 0.7728.

Практический вывод

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

• 0.6052 ± 0.0231 macro-F1 на 7 классах против 0.4971 ± 0.0185 у полного кадра;

• 0.7236 ± 0.0222 macro-F1 на 2 классах против 0.6752 ± 0.0266 у полного кадра.

Самый сильный общий вариант — ансамбль восьми входов:

• 0.7025 ± 0.0123 на 7 классах;

• 0.7728 ± 0.0303 на 2 классах.

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

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

Будем рады в комментариях почитать о ваших достижениях в похожих задачах, а также ответить на вопросы.

Ссылки

1. Набор данных MMOTU на Figshare: figshare.com/articles/dataset/_zip/25058690

2. MMOTU / DS2Net GitHub: github.com/cv516Buaa/MMOTU_DS2Net

3. Статья MMOTU: arxiv.org/abs/2207.06799

4. ConvNeXt: arxiv.org/abs/2201.03545

5. DenseNet: arxiv.org/abs/1608.06993

6. EfficientNetV2: arxiv.org/abs/2104.00298

7. macro-F1 и f1_score в scikit-learn: scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html

8. Оценка моделей в scikit-learn: scikit-learn.org/stable/modules/model_evaluation.html