Классификация инсульта по КТ-снимкам

Учебный проект по классификации КТ-снимков головного мозга на три класса: Normal, Bleeding, Ischemia.

Проект посвящён исследованию различных архитектур глубокого обучения и подходов к переносу знаний и тонкой настройки (transfer learning / fine-tuning) для классификации медицинских изображений.

В работе использовался исходный датасет Brain Stroke CT Dataset с Kaggle.

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

Задача

Для каждого КТ-снимка необходимо определить один из трёх классов:

• Normal — отсутствие инсульта;
• Bleeding — геморрагический инсульт;
• Ischemia — ишемический инсульт.

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

Исследованные модели

DenseNet121

Исследовано несколько вариантов использования предобученной DenseNet121 с весами ImageNet:

1. Transfer learning без fine-tuning — база DenseNet121 полностью заморожена, обучается только классификационная часть с дополнительным SE-блоком.
2. Fine-tuning последних слоёв — после обучения классификационной части размораживались последние 80 слоёв DenseNet121 и выполнялось дообучение с меньшей скоростью обучения.
3. Полная разморозка — все слои DenseNet121 становились обучаемыми для адаптации модели к медицинским изображениям.

Кастомная CNN

Разработана собственная компактная сверточная архитектура с остаточными блоками и SE-блоками.

Архитектура была ориентирована на работу с относительно небольшим объёмом медицинских данных и использовалась для проверки того, насколько собственная CNN может конкурировать с предобученной архитектурой.

LeNet-5

Реализована и адаптирована под задачу классификации КТ классическая архитектура LeNet-5.

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

Предобработка и обучение

В проекте реализован полный пайплайн подготовки изображений:

• чтение PNG;
• изменение размера изображений;
• нормализация;
• аугментация;
• работа с дисбалансом классов;
• формирование наборов train / validation / test;
• построение пайплайна загрузки данных с использованием tf.data.Dataset.

Для управления процессом обучения использовались:

• EarlyStopping;
• ReduceLROnPlateau;
• ModelCheckpoint.

Оценка моделей

Модели сравнивались по нескольким метрикам классификации:

• Accuracy;
• Precision;
• Recall;
• F1-score;
• Macro F1;
• Confusion Matrix;
• ROC-AUC.

Результаты

Лучший результат среди исследованных вариантов на тестовой выборке показала кастомная CNN:

• Accuracy = 0.705
• Macro F1 = 0.664
• Macro Recall = 0.755
• ROC-AUC ≈ 0.89–0.92

Структура проекта

Проект содержит три Jupyter Notebook, соответствующих основным направлениям экспериментов:

• эксперименты с DenseNet121 и различными режимами transfer learning / fine-tuning;
• эксперименты с кастомной CNN;
• эксперименты с LeNet-5, включая различные размеры входных изображений.

Связь с дипломным проектом

Полученный опыт был использован как основа для дальнейшей работы над дипломным проектом (https://github.com/aklnrn/brain-ct-stroke-analysis).

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