Удаление хирургического дыма: первые базовые модели

Задача

Электрокоагуляция заполняет поле зрения лапароскопа дымом. Задача – восстановить чистый кадр по задымлённому. Один и тот же момент нельзя снять и с дымом, и без него, поэтому мы учим модели на синтетических парах: дым, смоделированный в Blender, накладывается на чистые кадры открытого датасета CholecT50 с лапароскопическими холецистэктомиями.

• Обучение: 5 882 чистых кадра из 35 видео, дым каждую эпоху новый.
• Валидация: 518 фиксированных пар из 5 других видео.
• Разбиение: по видео, так что ни одна операция не попадает в обе части.

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

Результаты

Все три модели обучались в одной сессии на арендованном GPU A40, по 15 эпох на целых кадрах 854×480, с одним начальным состоянием (seed). В таблице – PSNR на валидационных парах в дБ, в скобках – разница с базовой линией «ничего не делать».

Residual U-Net предсказывает поправку к входному кадру, а не новое изображение. Он улучшает кадр при любой плотности дыма, сильнее всего – при среднем дыме.

Pix2Pix, генеративно-состязательная сеть (GAN), почти не меняет кадры со средним и плотным дымом и портит кадры с лёгким: его генератор рисует правдоподобную текстуру даже там, где кадр был почти чистым. Поэтому результат GAN нельзя считать тем, что скрывал дым.

Что не сработало и почему

• Вырезки. Сначала U-Net учился на вырезках 256×256 – это распространённая настройка по умолчанию.
  • Через 11 эпох он всё ещё уступал входу: 28,75 против 28,99 дБ.
  • Вырезки разбавляют дым и никогда не показывают тёмный край изображения эндоскопа.
  • На целых кадрах та же сеть обогнала вход на 2,4 дБ уже через 5 эпох.
• SRGAN, три попытки.
  1. Во время обучения выход обрезался до допустимого диапазона, и пиксели, вышедшие за него, теряли градиент, который мог бы их вернуть. За одну эпоху все пиксели стали чисто красными, PSNR составил 6,3 дБ.
  2. После исправления один прогон по ошибке запустился на старом коде и повторил сбой.
  3. В третьей попытке генератор сначала учится только на попиксельной ошибке, как в оригинальной статье, но всё равно ухудшал кадры по сравнению со входом: 25,4 дБ. Вероятнее всего, виноваты настройки оптимизатора, взятые от Pix2Pix; в статье использовались другие. На этом мы остановились, чтобы не тратить больше GPU-часов.
• Организация работы. Забытый процесс 2,5 часа обучался рядом с основным прогоном и вдвое замедлил его. Проверки на такие случаи теперь есть в инструкции по обучению.

Чем эти цифры не являются

• Это не законченная серия. Это результаты для одного начального состояния на валидационном разбиении. Улучшение мы заявляем только на минимум трёх начальных состояниях с непересекающимися доверительными интервалами, поэтому серия не завершена. Отложенную тестовую выборку мы не трогали.
• Это не настоящий дым. Дым синтетический, и ни один клиницист не сверял его с настоящим. Кадры с настоящим дымом пока не размечены.
• Это не полная оценка. Скорость, блики, кровь и расфокус не измерялись.
• Это не релиз. Ни одна модель не опубликована.

RESEARCH USE ONLY – NOT FOR CLINICAL USE (только для исследований, не для клинического применения). Восстановленные кадры нельзя показывать хирургу во время операции.

Возможные следующие шаги

• Ещё два начальных состояния для U-Net, затем одна зафиксированная проверка на тестовой выборке.
• Разметка кадров с настоящим дымом, чтобы проверять модели и на нём.
• Ещё одна попытка SRGAN с настройками оптимизатора из статьи.

Полная карточка модели с идентификаторами прогонов – в репозитории проекта. Смотрите также журнал экспериментов и заметку «Почему важны отрицательные результаты».