ИИ для анализа рентген‑снимков грудного отдела

Классификация шести патологий на рентгене: обучение только головы поверх RAD‑DINO‑MAIRA-2

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

Это не сегментация (модель не рисует маски органов) и не генерация текстового отчёта. Это многометочная классификация: на выходе — шесть чисел от 0 до 1. Например: «плевральный выпот — 87%, пневмоторакс — 12%, кардиомегалия — 3%» и так далее по остальным классам. Это не диагноз и не заключение врача, а вероятностная оценка по заранее заданному списку патологий.

Полное дообучение большой сети сразу не тянул: это дорого, долго, и на десятках тысяч снимков легко уйти в переобучение. Поэтому я заморозил тяжёлый предобученный кодировщик и обучал только небольшую голову. Macro AUC на проверочной выборке — 0.9025, примерно за десяток эпох. Обучение ^[1] шло на A100 80 ГБ — для этой схемы брать более мощную карту смысла не было. На RunPod ушло около $300.

Неудачных попыток было много — в основном на DenseNet121, прежде чем я перешёл на RAD‑DINO. Ниже — что в итоге сработало, где я ошибся с путями к данным и почему отказался от горизонтального отражения. Сейчас ищу организацию для клинических испытаний, чтобы регистрировать изделие 3 класса риска в РФ. Эта статья — про исследовательский прототип, а не про клиническое применение.

Дисклеймер. Это не СЭМД и не замена рентгенологу. В основе лежит microsoft/rad-dino-maira-2 под лицензией MSRLA — только для исследований.

Задача

Шесть классов из VinBigData Challenge:

Один снимок может одновременно относиться к нескольким классам. Поэтому на выходе стоят шесть независимых сигмоид, а не один Softmax.

DenseNet121 — почему отказался

Почти любой ИИ на вопрос «как классифицировать рентген» предлагает DenseNet121, предобученный на ImageNet. Я на это несколько раз повёлся. Обучение не раз рушилось именно на этой ветке.

Чтобы понять, куда сеть смотрит и на чём делает вывод, я строил тепловые карты (Grad‑CAM) поверх снимка: красные зоны — области, на которые модель опиралась сильнее всего. На скриншотах ниже — не итоговая RAD‑DINO, а именно прогоны DenseNet121 в Google Colab. Над каждым кадром три строки: GT (истинная метка), PRED (предсказание) и Conf (уверенность). Зелёные подписи — совпадение, красные — ошибка ^[2].

Ложные «здоров». В разметке — выпот, пневмоторакс или кардиомегалия, а DenseNet уверенно ставит «No Finding» с вероятностью 81–94%. Был случай с кардиомегалией при 99,3% уверенности, что находок нет. Карта при этом часто сидит в центре или средостении и не подсвечивает увеличенное сердце.

Перепутанный класс. GT — кардиомегалия, PRED — плевральный выпот с уверенностью 74%. Внимание ^[3] сидит не на сердце, а выше — в верхних долях и в области шеи.

Пропуск пневмоторакса. GT — пневмоторакс, PRED — “No Finding” с 51,7%. Heatmap концентрируется в центре, а не у края лёгкого, где обычно ищут воздух.

Кривое внимание даже когда модель «угадала». На части снимков с GT “No Finding” и верным PRED карта внимания липнет к плечам, шее и углам кадра, а не к лёгочным полям. В ноутбуке я оставлял себе пометки: «ищи странные, нелогичные или слишком мелкие пятна» и «проверь углы — что должно исчезнуть».

После таких прогонов DenseNet121 для меня закрыт. Модель зубрила артефакты кадра и давала красивый AUC на обучении при явных провалах на конкретных снимках. Часть этой беды — не только архитектура, а привычка стартовать с самого «популярного» грязного датасета: любая нейронка первым делом советует NIH ChestX‑ray, и на нём как раз легко получить эффект Ганса. Именно это подтолкнуло к RAD‑DINO‑MAIRA-2 и к более аккуратному выбору данных.

Карты внимания для итоговой RAD‑DINO в статью не выкладываю: здесь важно было показать, где ошибался DenseNet, а не сравнивать две визуализации бок о бок. Для отладки RAD‑DINO скрипт с attention map у меня тоже есть — он в конце статьи.

Как строилась тепловая карта DenseNet121

Коротко: берём последний свёрточный блок DenseNet, считаем Grad‑CAM для выбранного класса, растягиваем карту до размера снимка и накладываем поверх рентгена.

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import models, transforms


class DenseNetGradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.activations = None
        self.gradients = None
        target_layer.register_forward_hook(self._save_activation)
        target_layer.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module, inp, out):
        self.activations = out.detach()

    def _save_gradient(self, module, gin, gout):
        self.gradients = gout[0].detach()

    def __call__(self, x, class_idx=None):
        self.model.zero_grad(set_to_none=True)
        logits = self.model(x)
        if class_idx is None:
            class_idx = int(logits.argmax(dim=1).item())

        logits[0, class_idx].backward()
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = F.relu((weights * self.activations).sum(dim=1, keepdim=True))
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        return cam, logits.detach()


def overlay_heatmap(image_rgb, cam, alpha=0.45):
    h, w = image_rgb.shape[:2]
    heat = cv2.resize(cam, (w, h))
    color = cv2.applyColorMap(np.uint8(255 * heat), cv2.COLORMAP_JET)
    color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
    return np.clip((1 - alpha) * image_rgb + alpha * color, 0, 255).astype(np.uint8)


# пример использования
model = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1)
# ...подставьте свою голову / веса и число классов...
cam_engine = DenseNetGradCAM(model, model.features[-1])

raw = Image.open("xray.png").convert("RGB")
tfm = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
x = tfm(raw).unsqueeze(0)

cam, logits = cam_engine(x)
overlay = overlay_heatmap(np.array(raw), cam)
# дальше: сохранить картинку и подписать GT / PRED / Conf

Почему RAD‑DINO‑MAIRA-2

RAD‑DINO‑MAIRA-2 учили самообучением (DINOv2) примерно на 1,4 млн рентгенов: MIMIC‑CXR, CheXpert, NIH, PadChest, BRAX и ещё закрытый пул. Нормализация рассчитана под MIMIC‑CXR, вход — 518×518, патчи — 14×14. Для шести бинарных голов хватает тонкой головы поверх замороженного кодировщика: обучается меньше 0,25% параметров, на A100 80 ГБ пакет размером 64 идёт спокойно.

На VinBigData снимки относительно чистые, поэтому случайное затирание кусков кадра в финал я не ставил. На сырых DICOM из реанимации без такой аугментации я бы не полез.

В разработке ИИ помогал с архитектурными решениями (в основном Claude Sonnet). Gemini тоже пробовал — он чаще ошибался и быстрее терял контекст длинной сессии. Код всё равно проверял сам.

Главное узкое место — не архитектура, а данные

Самое важное для такого проекта — чистые данные, а не «ещё один слой» или другой backbone. Архитектуру можно заменить за неделю. Грязную разметку вы будете тащить за собой месяцами: модель выучит шум, AUC на hold‑out из той же каши останется красивым, а на реальных снимках всё развалится. Для обучения не хватает именно чистых и согласованных меток — много снимков без радиологической проверки, с конфликтами между разметчиками и с артефактами, которые коррелируют с классом.

Отдельное предупреждение: не берите для обучения самый популярный датасет NIH ChestX‑ray / ChestX‑ray14 только потому, что он «большой и все его советуют». Это как раз тот набор, который любая нейронка и почти любой туториал рекомендуют первым. Проблема в том, что он слишком грязный: метки часто извлечены из текстов отчётов автоматически, много шума, ложных корреляций и «коротких путей». DenseNet на нём легко показывает приличный AUC и при этом смотрит не туда — ровно то, что я ловил на тепловых картах.

Большой грязный датасет хуже небольшого чистого. Если ставите эксперимент с нуля, лучше VinBigData / CheXpert / MIMIC‑CXR с понятной политикой разметки и ручной проверкой хотя бы части ошибок, чем гнать эпохи на ChestX‑ray14 “потому что так написал чат‑бот”.

Данные — всё с Kaggle

Качал только полные наборы с Kaggle. Если у источника не было готовых снимков в архиве — такой датасет не брал. ChestX‑ray14 в финальный пайплайн не брал именно из‑за качества меток, а не из‑за размера.

VinBigData

Использовал vinbigdata-512-image-dataset: 15 000 снимков 512×512 с разметкой радиологов. Объём около 8 ГБ, на RunPod загружал через Kaggle API. Это основа пайплайна.

VIN_CLASS_MAP = {3: 0, 4: 1, 5: 2, 10: 3, 12: 4, 6: 5}

Около часа убил на путь к файлам: в туториалах часто фигурирует /data/vin/train/, а в архиве awsaf49 лежит vinbigdata/train/. Если папки нет — получаешь 0 снимков и уходишь в отладку через os.path.exists.

OpenI

Датасет chest-xrays-indiana-university тоже взял с Kaggle, потому что снимки шли в комплекте. В финальном прогоне OpenI использовал. Это 7466 кадров. Готовых меток под наши шесть классов нет — только текстовые отчёты. Написал простой разметчик на правилах:

RULES = {
    0: ["cardiomegaly", "cardiac enlargement", "heart is enlarged"],
    1: ["atelectasis", "atelectatic", "collapse", "collapsed"],
    2: ["consolidation", "consolidative", "opacification"],
    3: ["pleural effusion", "effusion", "pleural fluid"],
    4: ["pneumothorax"],
    5: ["pulmonary edema", "edema", "vascular congestion"],
}

NEGATIONS = [
    "no ",
    "not ",
    "without ",
    "negative for ",
    "no evidence of ",
]

Разметка шумная: слово effusion иногда цепляет лишнее. Зато это дало дополнительный объём данных.

Итого

В сумме получилось 22 466 снимка. Деление 85/15: обучение — 19 096, проверка — 3 370.

pos_weight шёл в BCEWithLogitsLoss. Рентгенологов на разбор ошибок не звал — смотрел метрики и карты внимания сам.

Архитектура

Пайплайн такой:

1. вход 518×518

2. AutoImageProcessor (нормализация под MIMIC‑CXR)

3. RAD‑DINO‑MAIRA-2 — заморожен, ~86,5 млн параметров

4. CLS‑токен размерности 768

5. голова: LayerNorm → Linear(768→256) → GELU → Dropout(0.3) → Linear(256→6)

6. шесть логитов → сигмоида

class ChestClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(
            "microsoft/rad-dino-maira-2",
            trust_remote_code=True,
        )
        for p in self.backbone.parameters():
            p.requires_grad = False

        self.head = nn.Sequential(
            nn.LayerNorm(768),
            nn.Linear(768, 256),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(256, 6),
        )

    def forward(self, pixel_values):
        cls = self.backbone(pixel_values=pixel_values).last_hidden_state[:, 0, :]
        return self.head(cls)

Обучается 199 942 параметра из 86 780 422.

Аугментации

train_aug = A.Compose([
    A.RandomResizedCrop(
        size=(518, 518),
        scale=(0.85, 1.0),
        ratio=(0.95, 1.05),
    ),
    A.ShiftScaleRotate(
        shift_limit=0.05,
        scale_limit=0,
        rotate_limit=7,
        p=0.5,
    ),
    A.RandomBrightnessContrast(
        brightness_limit=0.1,
        contrast_limit=0.1,
        p=0.5,
    ),
])

Горизонтальное отражение не ставил: сердце слева, зеркальное отражение ломает анатомию.

Два набора — два веса в функции потерь

LOSS_W = torch.tensor([3.0, 2.0])  # vin, openi


def loss_fn(logits, labels, src_ids):
    base = F.binary_cross_entropy_with_logits(
        logits, labels, pos_weight=pos_weight, reduction="none"
    )
    return (base * LOSS_W[src_ids].unsqueeze(1)).mean()

Плюс WeightedRandomSampler: OpenI — 5×, Vin — 1×. Без этого редкие классы проседали.

Обучение

Обучал на RunPod, A100 80 ГБ. Бюджет этой ветки — около $300.

Урок после кучи сорванных прогонов: после каждой эпохи сохранять лучшие веса и сразу закладывать выход. Если модель N эпох подряд себя не превосходит — стоп. Без этого легко полдня жечь GPU впустую.

Epoch 00 | loss=2.616 | AUC=0.880
Epoch 03 | loss=1.939 | AUC=0.899
Epoch 09 | loss=1.621 | AUC=0.903  ← лучший
Epoch 16 | ранняя остановка, best AUC=0.9025

После 9-й эпохи началось плато. Полное дообучение основы не гонял.

Метрики по классам

Лучший чекпоинт на RAD‑DINO:

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

Прогон и карта внимания для RAD‑DINO

Для итоговой модели оставил скрипт, который считает вероятности и строит карту внимания последнего слоя трансформера. Это не Grad‑CAM — просто видно, куда смотрит блок. На DenseNet так я и ловил зубрёжку артефактов. На RAD‑DINO в финальном прогоне картами пользовался реже, в статью их не выкладывал, но код оставил для отладки.

processor = AutoImageProcessor.from_pretrained("microsoft/rad-dino-maira-2")

model = ChestClassifier()
model.load_state_dict(torch.load("best_model.pth"))
model.eval()

raw_img = Image.open(path).convert("RGB")
inputs = processor(images=raw_img, return_tensors="pt")

with torch.no_grad():
    outputs = model.backbone(**inputs, output_attentions=True)
    logits = model.head(outputs.last_hidden_state[:, 0, :])
    probs = torch.sigmoid(logits[0])

    attn = outputs.attentions[-1][0][:, 0, 1:].mean(dim=0)
    grid = int(np.sqrt(attn.shape[0]))
    heatmap = attn.reshape(grid, grid).numpy()

Голова весит около 800 КБ. Основу качает Hugging Face под MSRLA. Веса и этапы у меня локально / в закрытом git — публично не выкладывал.

Что дальше

Разморозил бы последние 2–4 блока кодировщика. OpenI почистил бы жёстче или вовсе выкинул. На реанимационных снимках добавил бы затирание областей и обрезку трубок. Внешняя проверка обязательна. Для итоговой модели нужен такой же разбор ошибок по снимкам, как я делал для DenseNet.

Сейчас ищу организацию под клинические испытания для регистрации изделия 3 класса риска в РФ.

По цифрам

Октябрь 2025 — апрель 2026. Данные — с Kaggle, только полные наборы со снимками. Железо — A100 80 ГБ, около $300. Обучал только голову поверх RAD‑DINO‑MAIRA-2. Macro AUC — 0.90. Для исследовательского прототипа этого хватает. Для клиники — пока нет.

Статья про исследование, а не про зарегистрированное изделие.