{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "0", "metadata": {}, "outputs": [], "source": [ "%pip install nibabel monai" ] }, { "cell_type": "raw", "id": "1", "metadata": {}, "source": [ "Task09_Spleen (селезёнка) :\n", "\n", "Размер: всего ~1.5 ГБ (скачается быстро)\n", "\n", "Структура: простая бинарная сегментация (есть селезёнка / нет селезёнки)\n", "\n", "Количество файлов: 41 тренировочное + 20 тестовых изображений" ] }, { "cell_type": "code", "execution_count": null, "id": "2", "metadata": {}, "outputs": [], "source": [ "# Task09_Spleen (рекомендую для начала) - 1.5 ГБ\n", "!wget -O Task09_Spleen.tar 'https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar'" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "import tarfile\n", "import os\n", "import nibabel as nib\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "# Распаковка\n", "with tarfile.open('Task09_Spleen.tar', 'r') as tar:\n", " tar.extractall()\n", "\n", "# Посмотрим содержимое\n", "base_dir = 'Task09_Spleen'\n", "print(\"Содержимое:\", os.listdir(base_dir))\n", "# Ожидается: imagesTr (тренировочные изображения), labelsTr (маски), imagesTs (тестовые)" ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "base_dir = 'Task09_Spleen'\n", "print(\"Папки в датасете:\", os.listdir(base_dir))\n", "# Ожидаем: ['imagesTr', 'labelsTr', 'imagesTs', 'dataset.json']" ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [], "source": [ "train_images = sorted(os.listdir(os.path.join(base_dir, 'imagesTr')))\n", "train_labels = sorted(os.listdir(os.path.join(base_dir, 'labelsTr')))\n", "\n", "print(f\"Тренировочных изображений: {len(train_images)}\")\n", "print(f\"Масок: {len(train_labels)}\")\n", "print(f\"Пример имени файла изображения: {train_images[0]}\")\n", "print(f\"Пример имени файла маски: {train_labels[0]}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": { "tags": [] }, "outputs": [], "source": [ "import os\n", "import nibabel as nib\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "base_dir = 'Task09_Spleen'\n", "\n", "# Получаем списки файлов, исключая служебные (начинающиеся с ._)\n", "train_images = sorted([\n", " f for f in os.listdir(os.path.join(base_dir, 'imagesTr')) \n", " if not f.startswith('._') and f.endswith('.nii.gz')\n", "])\n", "train_labels = sorted([\n", " f for f in os.listdir(os.path.join(base_dir, 'labelsTr')) \n", " if not f.startswith('._') and f.endswith('.nii.gz')\n", "])\n", "\n", "print(f\"Тренировочных изображений (реальных): {len(train_images)}\")\n", "print(f\"Масок (реальных): {len(train_labels)}\")\n", "print(f\"Пример имени файла изображения: {train_images[0]}\")\n", "print(f\"Пример имени файла маски: {train_labels[0]}\")\n", "\n", "# Проверим, совпадают ли имена (в MSD они могут немного отличаться, например, у изображений может быть суффикс '_0000')\n", "# Для простоты возьмём первый файл из изображений и попробуем найти соответствующую маску\n", "# Обычно в Task09_Spleen имена совпадают (без суффиксов)" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "# Выбираем первый реальный файл изображения и соответствующую маску\n", "# Предполагаем, что имена совпадают\n", "img_name = train_images[0]\n", "# Ищем маску с таким же именем (в Task09_Spleen они одинаковые)\n", "mask_name = img_name # если имена совпадают\n", "\n", "img_path = os.path.join(base_dir, 'imagesTr', img_name)\n", "mask_path = os.path.join(base_dir, 'labelsTr', mask_name)\n", "\n", "# Загрузка\n", "img_nii = nib.load(img_path)\n", "mask_nii = nib.load(mask_path)\n", "\n", "img_data = img_nii.get_fdata()\n", "mask_data = mask_nii.get_fdata()\n", "\n", "print(f\"Размер изображения: {img_data.shape}\")\n", "print(f\"Уникальные значения в маске: {np.unique(mask_data)}\") # должно быть [0., 1.] (фон и селезёнка)\n", "\n", "# Покажем несколько срезов\n", "slices = np.linspace(0, img_data.shape[2]-1, 5, dtype=int)\n", "\n", "plt.figure(figsize=(15, 10))\n", "for i, idx in enumerate(slices):\n", " plt.subplot(2, 5, i+1)\n", " plt.imshow(img_data[:, :, idx], cmap='gray')\n", " plt.title(f'Срез {idx}')\n", " plt.axis('off')\n", "\n", " plt.subplot(2, 5, i+6)\n", " plt.imshow(mask_data[:, :, idx], cmap='hot', alpha=0.7)\n", " plt.imshow(img_data[:, :, idx], cmap='gray', alpha=0.3)\n", " plt.title(f'Маска срез {idx}')\n", " plt.axis('off')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": { "tags": [] }, "outputs": [], "source": [ "import os\n", "import torch\n", "import nibabel as nib\n", "import numpy as np\n", "from monai.networks.nets import UNet\n", "from monai.losses import DiceCELoss\n", "from monai.metrics import DiceMetric\n", "from monai.inferers import sliding_window_inference\n", "from monai.transforms import (\n", " Compose, LoadImaged, EnsureChannelFirstd, ScaleIntensityd,\n", " RandRotate90d, RandFlipd, RandShiftIntensityd, ToTensord,\n", " SpatialPadd\n", ")\n", "from monai.data import CacheDataset, DataLoader\n", "\n", "# --- 1. Подготовка данных ---\n", "base_dir = 'Task09_Spleen'\n", "images_dir = os.path.join(base_dir, 'imagesTr')\n", "labels_dir = os.path.join(base_dir, 'labelsTr')\n", "\n", "train_images = sorted([os.path.join(images_dir, f) \n", " for f in os.listdir(images_dir) \n", " if f.endswith('.nii.gz') and not f.startswith('._')])\n", "train_labels = sorted([os.path.join(labels_dir, f) \n", " for f in os.listdir(labels_dir) \n", " if f.endswith('.nii.gz') and not f.startswith('._')])\n", "\n", "assert len(train_images) == len(train_labels), \"Количество изображений и масок не совпадает\"\n", "print(f\"Найдено пар изображение-маска: {len(train_images)}\")\n", "\n", "# --- 2. Определение максимального размера и округление до кратного 16 ---\n", "print(\"Определяем максимальный размер по всем осям...\")\n", "max_shape = [0, 0, 0]\n", "for img_path in train_images:\n", " img = nib.load(img_path)\n", " shape = img.shape\n", " max_shape = [max(max_shape[i], shape[i]) for i in range(3)]\n", "print(f\"Максимальный размер (H, W, D): {max_shape}\")\n", "\n", "# Округляем до ближайшего числа, кратного 16\n", "aligned_shape = [((dim + 15) // 16) * 16 for dim in max_shape]\n", "print(f\"Размер после выравнивания (кратно 16): {aligned_shape}\")\n", "\n", "# --- 3. Создание словарей данных и разбиение на train/val ---\n", "data_dicts = [{\"image\": img, \"label\": lbl} for img, lbl in zip(train_images, train_labels)]\n", "split_idx = int(0.8 * len(data_dicts))\n", "train_dicts = data_dicts[:split_idx]\n", "val_dicts = data_dicts[split_idx:]\n", "print(f\"Тренировочных: {len(train_dicts)}, Валидационных: {len(val_dicts)}\")\n", "\n", "# --- 4. Трансформы (используем aligned_shape) ---\n", "train_transforms = Compose([\n", " LoadImaged(keys=[\"image\", \"label\"]),\n", " EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n", " ScaleIntensityd(keys=[\"image\"]),\n", " SpatialPadd(keys=[\"image\", \"label\"], spatial_size=aligned_shape, method=\"symmetric\"),\n", " RandRotate90d(keys=[\"image\", \"label\"], prob=0.5, spatial_axes=(0, 1)),\n", " RandFlipd(keys=[\"image\", \"label\"], prob=0.5, spatial_axis=0),\n", " RandShiftIntensityd(keys=[\"image\"], offsets=0.1, prob=0.5),\n", " ToTensord(keys=[\"image\", \"label\"])\n", "])\n", "\n", "val_transforms = Compose([\n", " LoadImaged(keys=[\"image\", \"label\"]),\n", " EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n", " ScaleIntensityd(keys=[\"image\"]),\n", " SpatialPadd(keys=[\"image\", \"label\"], spatial_size=aligned_shape, method=\"symmetric\"),\n", " ToTensord(keys=[\"image\", \"label\"])\n", "])\n", "\n", "# --- 5. Датасеты и загрузчики ---\n", "train_ds = CacheDataset(data=train_dicts, transform=train_transforms, cache_rate=0.5)\n", "val_ds = CacheDataset(data=val_dicts, transform=val_transforms, cache_rate=0.5)\n", "\n", "batch_size = 2\n", "train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)\n", "val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2)\n", "\n", "# --- 6. Модель, функция потерь, оптимизатор ---\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(f\"Используется устройство: {device}\")\n", "\n", "model = UNet(\n", " spatial_dims=3,\n", " in_channels=1,\n", " out_channels=2,\n", " channels=(16, 32, 64, 128, 256),\n", " strides=(2, 2, 2, 2),\n", " num_res_units=2,\n", ").to(device)\n", "\n", "loss_function = DiceCELoss(to_onehot_y=True, softmax=True)\n", "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n", "dice_metric = DiceMetric(include_background=False, reduction=\"mean\")\n", "\n", "roi_size = (96, 96, 96)\n", "sw_batch_size = 4\n", "\n", "# --- 7. Цикл обучения ---\n", "max_epochs = 50\n", "val_interval = 2\n", "best_metric = -1\n", "best_metric_epoch = -1\n", "\n", "for epoch in range(max_epochs):\n", " print(f\"Эпоха {epoch+1}/{max_epochs}\")\n", " model.train()\n", " epoch_loss = 0\n", " step = 0\n", " for batch_data in train_loader:\n", " inputs, labels = batch_data[\"image\"].to(device), batch_data[\"label\"].to(device)\n", " optimizer.zero_grad()\n", " outputs = model(inputs)\n", " loss = loss_function(outputs, labels)\n", " loss.backward()\n", " optimizer.step()\n", " epoch_loss += loss.item()\n", " step += 1\n", " \n", " epoch_loss /= step\n", " print(f\" Средняя потеря: {epoch_loss:.4f}\")\n", "\n", " if (epoch + 1) % val_interval == 0:\n", " model.eval()\n", " with torch.no_grad():\n", " for val_data in val_loader:\n", " val_inputs, val_labels = val_data[\"image\"].to(device), val_data[\"label\"].to(device)\n", " val_pred = sliding_window_inference(\n", " val_inputs, roi_size, sw_batch_size, model, overlap=0.5\n", " )\n", " dice_metric(y_pred=val_pred, y=val_labels)\n", " metric = dice_metric.aggregate().item()\n", " dice_metric.reset()\n", " print(f\" Dice на валидации: {metric:.4f}\")\n", " if metric > best_metric:\n", " best_metric = metric\n", " best_metric_epoch = epoch + 1\n", " torch.save(model.state_dict(), \"best_spleen_model.pth\")\n", " print(\" ✅ Модель сохранена\")\n", "\n", "print(f\"🎉 Лучший Dice: {best_metric:.4f} на эпохе {best_metric_epoch}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from monai.networks.nets import UNet\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "\n", "# Создаём архитектуру (те же параметры, что и при обучении)\n", "model = UNet(\n", " spatial_dims=3,\n", " in_channels=1,\n", " out_channels=2,\n", " channels=(16, 32, 64, 128, 256),\n", " strides=(2, 2, 2, 2),\n", " num_res_units=2,\n", ").to(device)\n", "\n", "# Загружаем веса\n", "checkpoint = torch.load(\"best_spleen_model.pth\", map_location=device)\n", "model.load_state_dict(checkpoint)\n", "model.eval()\n", "print(\"✅ Модель загружена и готова к инференсу.\")" ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "# Параметры (те же, что использовались при обучении)\n", "target_shape = [512, 512, 176] # выровненный размер\n", "\n", "# Трансформы для инференса\n", "infer_transform = Compose([\n", " LoadImage(image_only=True),\n", " EnsureChannelFirst(),\n", " ScaleIntensity(),\n", " SpatialPad(spatial_size=target_shape, method=\"symmetric\"),\n", " ToTensor()\n", "])\n", "\n", "# Выберите тестовый файл (например, первый из imagesTs)\n", "import os\n", "test_dir = 'Task09_Spleen/imagesTs'\n", "if os.path.exists(test_dir):\n", " test_files = [f for f in os.listdir(test_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " if test_files:\n", " test_path = os.path.join(test_dir, test_files[0])\n", " else:\n", " # Если нет тестовых, возьмём валидационный из тренировочных (например, 35-й)\n", " train_dir = 'Task09_Spleen/imagesTr'\n", " train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " test_path = os.path.join(train_dir, train_files[-5]) # последний файл\n", "else:\n", " train_dir = 'Task09_Spleen/imagesTr'\n", " train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " test_path = os.path.join(train_dir, train_files[-5])\n", "\n", "print(f\"Тестовый файл: {test_path}\")\n", "\n", "# Применяем трансформации и добавляем batch dimension\n", "input_tensor = infer_transform(test_path).unsqueeze(0).to(device) # [1,1,H,W,D]" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "from monai.inferers import sliding_window_inference\n", "\n", "roi_size = (96, 96, 96) # размер окна (можно подобрать)\n", "sw_batch_size = 4\n", "\n", "with torch.no_grad():\n", " # Прямой проход (если изображение целиком влезает, можно использовать model(input_tensor))\n", " # Но для 3D лучше sliding window:\n", " pred_logits = sliding_window_inference(\n", " input_tensor, roi_size, sw_batch_size, model, overlap=0.5\n", " )\n", " pred_mask = torch.argmax(pred_logits, dim=1).cpu().numpy()[0] # [H,W,D]\n", "\n", "print(\"Предсказание выполнено.\")\n", "print(f\"Уникальные значения в маске: {np.unique(pred_mask)}\") # должно быть [0,1]" ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "# Загружаем оригинальное изображение (без трансформаций)\n", "import nibabel as nib\n", "orig_img_nii = nib.load(test_path)\n", "orig_img = orig_img_nii.get_fdata() # [H,W,D] исходный размер\n", "\n", "# Размеры\n", "H, W, D = orig_img.shape\n", "print(f\"Исходный размер: {orig_img.shape}\")\n", "\n", "# Обрежем предсказанную маску до исходного размера (если нужно)\n", "pred_mask_cropped = pred_mask[:H, :W, :D]\n", "\n", "# Выберем несколько срезов по оси Z\n", "slices = np.linspace(0, D-1, 5, dtype=int)\n", "\n", "plt.figure(figsize=(15, 10))\n", "for i, idx in enumerate(slices):\n", " plt.subplot(2, 5, i+1)\n", " plt.imshow(orig_img[:, :, idx], cmap='gray')\n", " plt.title(f'Срез {idx}')\n", " plt.axis('off')\n", "\n", " plt.subplot(2, 5, i+6)\n", " plt.imshow(pred_mask_cropped[:, :, idx], cmap='hot', alpha=0.7)\n", " plt.imshow(orig_img[:, :, idx], cmap='gray', alpha=0.3)\n", " plt.title(f'Маска срез {idx}')\n", " plt.axis('off')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "from monai.networks.nets import UNet\n", "from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor\n", "from monai.inferers import sliding_window_inference\n", "from monai.metrics import DiceMetric, HausdorffDistanceMetric\n", "import os\n", "import nibabel as nib\n", "\n", "# 1. Загружаем модель (как вы уже сделали)\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model = UNet(\n", " spatial_dims=3,\n", " in_channels=1,\n", " out_channels=2,\n", " channels=(16, 32, 64, 128, 256),\n", " strides=(2, 2, 2, 2),\n", " num_res_units=2,\n", ").to(device)\n", "checkpoint = torch.load(\"best_spleen_model.pth\", map_location=device)\n", "model.load_state_dict(checkpoint)\n", "model.eval()\n", "\n", "# 2. Параметры (те же, что при обучении)\n", "target_shape = [512, 512, 176] # выровненный размер\n", "roi_size = (96, 96, 96)\n", "sw_batch_size = 4\n", "\n", "# Трансформы для инференса\n", "infer_transform = Compose([\n", " LoadImage(image_only=True),\n", " EnsureChannelFirst(),\n", " ScaleIntensity(),\n", " SpatialPad(spatial_size=target_shape, method=\"symmetric\"),\n", " ToTensor()\n", "])\n", "\n", "# 3. Собираем все валидационные пары (изображение + маска)\n", "base_dir = 'Task09_Spleen'\n", "val_images = sorted([os.path.join(base_dir, 'imagesTr', f) \n", " for f in os.listdir(os.path.join(base_dir, 'imagesTr')) \n", " if f.endswith('.nii.gz') and not f.startswith('._')])\n", "val_labels = sorted([os.path.join(base_dir, 'labelsTr', f) \n", " for f in os.listdir(os.path.join(base_dir, 'labelsTr')) \n", " if f.endswith('.nii.gz') and not f.startswith('._')])\n", "\n", "# Предположим, что последние 9 файлов были валидационными (как в вашем разбиении)\n", "val_images = val_images[-9:]\n", "val_labels = val_labels[-9:]\n", "\n", "print(f\"Валидационных образцов: {len(val_images)}\")\n", "\n", "# 4. Метрики\n", "dice_metric = DiceMetric(include_background=False, reduction=\"mean\")\n", "hausdorff_metric = HausdorffDistanceMetric(include_background=False, percentile=95)\n", "\n", "# 5. Цикл по всем валидационным образцам\n", "all_dice = []\n", "all_hd95 = []\n", "\n", "with torch.no_grad():\n", " for img_path, mask_path in zip(val_images, val_labels):\n", " # Подготовка изображения\n", " input_tensor = infer_transform(img_path).unsqueeze(0).to(device)\n", " \n", " # Предсказание\n", " pred_logits = sliding_window_inference(\n", " input_tensor, roi_size, sw_batch_size, model, overlap=0.5\n", " )\n", " pred_mask = torch.argmax(pred_logits, dim=1).cpu().numpy()[0] # [H,W,D]\n", " \n", " # Загрузка истинной маски (без паддинга)\n", " true_mask_nii = nib.load(mask_path)\n", " true_mask = true_mask_nii.get_fdata()\n", " \n", " # Обрезаем предсказание до исходного размера\n", " H, W, D = true_mask.shape\n", " pred_mask_cropped = pred_mask[:H, :W, :D]\n", " \n", " # Преобразуем в тензоры для MONAI метрик\n", " pred_tensor = torch.from_numpy(pred_mask_cropped).unsqueeze(0).unsqueeze(0).long() # [1,1,H,W,D]\n", " true_tensor = torch.from_numpy(true_mask).unsqueeze(0).unsqueeze(0).long()\n", " \n", " # Вычисляем Dice\n", " dice = dice_metric(pred_tensor, true_tensor)\n", " all_dice.append(dice.item())\n", " \n", " # Вычисляем Hausdorff Distance 95%\n", " hd = hausdorff_metric(pred_tensor, true_tensor)\n", " all_hd95.append(hd.item())\n", " \n", " print(f\"Файл: {os.path.basename(img_path)} -> Dice: {dice.item():.4f}, HD95: {hd.item():.2f} мм\")\n", "\n", "# 6. Итоговая статистика\n", "print(\"\\n\" + \"=\"*40)\n", "print(f\"Средний Dice по валидации: {np.mean(all_dice):.4f} ± {np.std(all_dice):.4f}\")\n", "print(f\"Средний HD95 по валидации: {np.mean(all_hd95):.2f} ± {np.std(all_hd95):.2f} мм\")\n", "print(\"=\"*40)" ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": { "tags": [] }, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "from monai.networks.nets import UNet\n", "from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor\n", "from monai.inferers import sliding_window_inference\n", "import nibabel as nib\n", "import os\n", "\n", "# --- 1. Загрузка модели ---\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model = UNet(\n", " spatial_dims=3,\n", " in_channels=1,\n", " out_channels=2,\n", " channels=(16, 32, 64, 128, 256),\n", " strides=(2, 2, 2, 2),\n", " num_res_units=2,\n", ").to(device)\n", "checkpoint = torch.load(\"best_spleen_model.pth\", map_location=device)\n", "model.load_state_dict(checkpoint)\n", "model.eval()\n", "print(\"✅ Модель загружена\")\n", "\n", "# --- 2. Вывод структуры model.model для выбора целевого слоя ---\n", "print(\"\\nСтруктура model.model:\")\n", "for idx, module in enumerate(model.model):\n", " print(f\"[{idx}]: {module.__class__.__name__}\")\n", "\n", "# --- 3. Выбор целевого слоя ---\n", "# Обычно для Grad‑CAM берут последний сверточный слой энкодера.\n", "# В U‑Net от MONAI энкодер состоит из блоков `Down`.\n", "# Выберите индекс блока, который вы хотите использовать (например, последний Down перед bottleneck).\n", "# Попробуйте индекс 2 или 3, если они есть.\n", "# В вашем случае, возможно, подойдёт индекс 2.\n", "# Если не уверены, начните с последнего Down блока.\n", "target_layer_idx = 2 # Попробуйте изменить на 2 или 3, в зависимости от вывода выше\n", "target_layer = model.model[target_layer_idx]\n", "print(f\"\\nВыбран целевой слой: model.model[{target_layer_idx}] -> {target_layer.__class__.__name__}\")\n", "\n", "# --- 4. Класс Grad-CAM ---\n", "class GradCAM:\n", " def __init__(self, model, target_layer):\n", " self.model = model\n", " self.target_layer = target_layer\n", " self.gradients = None\n", " self.activations = None\n", " \n", " target_layer.register_forward_hook(self.save_activation)\n", " target_layer.register_full_backward_hook(self.save_gradient) # обратите внимание на register_full_backward_hook\n", " \n", " def save_activation(self, module, input, output):\n", " self.activations = output.detach()\n", " \n", " def save_gradient(self, module, grad_input, grad_output):\n", " self.gradients = grad_output[0].detach()\n", " \n", " def generate_cam(self, input_tensor, target_class=1):\n", " self.model.zero_grad()\n", " output = self.model(input_tensor)\n", " \n", " target = output[:, target_class].sum()\n", " target.backward()\n", " \n", " weights = torch.mean(self.gradients, dim=(2, 3, 4), keepdim=True)\n", " cam = torch.sum(weights * self.activations, dim=1, keepdim=True)\n", " cam = torch.relu(cam)\n", " cam = cam / (cam.max() + 1e-8)\n", " return cam.squeeze().cpu().numpy()\n", "\n", "# --- 5. Подготовка тестового изображения ---\n", "test_dir = 'Task09_Spleen/imagesTs'\n", "if os.path.exists(test_dir):\n", " test_files = [f for f in os.listdir(test_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " if test_files:\n", " test_path = os.path.join(test_dir, test_files[0])\n", " else:\n", " train_dir = 'Task09_Spleen/imagesTr'\n", " train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " test_path = os.path.join(train_dir, train_files[-5])\n", "else:\n", " train_dir = 'Task09_Spleen/imagesTr'\n", " train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]\n", " test_path = os.path.join(train_dir, train_files[-5])\n", "\n", "print(f\"\\nТестовый файл: {test_path}\")\n", "\n", "target_shape = [512, 512, 176]\n", "infer_transform = Compose([\n", " LoadImage(image_only=True),\n", " EnsureChannelFirst(),\n", " ScaleIntensity(),\n", " SpatialPad(spatial_size=target_shape, method=\"symmetric\"),\n", " ToTensor()\n", "])\n", "\n", "input_tensor = infer_transform(test_path).unsqueeze(0).to(device)\n", "\n", "# --- 6. Получение Grad-CAM ---\n", "grad_cam = GradCAM(model, target_layer)\n", "cam_3d = grad_cam.generate_cam(input_tensor, target_class=1)\n", "\n", "# --- 7. Визуализация ---\n", "orig_img = nib.load(test_path).get_fdata()\n", "H, W, D = orig_img.shape\n", "\n", "slices = np.linspace(0, D-1, 5, dtype=int)\n", "\n", "plt.figure(figsize=(15, 10))\n", "for i, idx in enumerate(slices):\n", " cam_slice = cam_3d[:H, :W, idx]\n", " \n", " plt.subplot(2, 5, i+1)\n", " plt.imshow(orig_img[:, :, idx], cmap='gray')\n", " plt.title(f'Срез {idx}')\n", " plt.axis('off')\n", " \n", " plt.subplot(2, 5, i+6)\n", " plt.imshow(orig_img[:, :, idx], cmap='gray', alpha=0.5)\n", " plt.imshow(cam_slice, cmap='jet', alpha=0.5, interpolation='bilinear')\n", " plt.title(f'Grad-CAM срез {idx}')\n", " plt.axis('off')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": { "tags": [] }, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "from monai.networks.nets import UNet\n", "from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor\n", "import nibabel as nib\n", "import os\n", "\n", "# --- 1. Загрузка модели ---\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model = UNet(\n", " spatial_dims=3,\n", " in_channels=1,\n", " out_channels=2,\n", " channels=(16, 32, 64, 128, 256),\n", " strides=(2, 2, 2, 2),\n", " num_res_units=2,\n", ").to(device)\n", "checkpoint = torch.load(\"best_spleen_model.pth\", map_location=device)\n", "model.load_state_dict(checkpoint)\n", "model.eval()\n", "print(\"✅ Модель загружена\")\n", "\n", "# --- 2. Вывод структуры model.model ---\n", "print(\"\\n📋 Доступные слои в model.model:\")\n", "for idx, module in enumerate(model.model):\n", " print(f\" [{idx}]: {module.__class__.__name__}\")\n", "\n", "# --- 3. Класс Grad-CAM ---\n", "class GradCAM:\n", " def __init__(self, model, target_layer):\n", " self.model = model\n", " self.target_layer = target_layer\n", " self.gradients = None\n", " self.activations = None\n", " \n", " target_layer.register_forward_hook(self.save_activation)\n", " target_layer.register_full_backward_hook(self.save_gradient)\n", " \n", " def save_activation(self, module, input, output):\n", " self.activations = output.detach()\n", " \n", " def save_gradient(self, module, grad_input, grad_output):\n", " self.gradients = grad_output[0].detach()\n", " \n", " def generate_cam(self, input_tensor, target_class=1):\n", " self.model.zero_grad()\n", " output = self.model(input_tensor)\n", " \n", " target = output[:, target_class].sum()\n", " target.backward()\n", " \n", " weights = torch.mean(self.gradients, dim=(2, 3, 4), keepdim=True)\n", " cam = torch.sum(weights * self.activations, dim=1, keepdim=True)\n", " cam = torch.relu(cam)\n", " cam = cam / (cam.max() + 1e-8)\n", " return cam.squeeze().cpu().numpy() # [H, W, D]\n", "\n", "# --- 4. Подготовка тестовых данных ---\n", "base_dir = 'Task09_Spleen'\n", "img_dir = os.path.join(base_dir, 'imagesTr') # используем тренировочные, если нет test\n", "label_dir = os.path.join(base_dir, 'labelsTr')\n", "\n", "# Получаем списки файлов\n", "img_files = sorted([f for f in os.listdir(img_dir) if f.endswith('.nii.gz') and not f.startswith('._')])\n", "label_files = sorted([f for f in os.listdir(label_dir) if f.endswith('.nii.gz') and not f.startswith('._')])\n", "\n", "print(f\"\\n📂 Найдено изображений: {len(img_files)}, масок: {len(label_files)}\")\n", "\n", "# Создаём словарь соответствий: имя маски (без пути) -> полный путь\n", "label_dict = {os.path.basename(f): os.path.join(label_dir, f) for f in label_files}\n", "\n", "# Выбираем файл изображения (например, 10-й)\n", "idx_file = 10 # можно изменить\n", "img_filename = img_files[idx_file]\n", "img_path = os.path.join(img_dir, img_filename)\n", "\n", "# Ищем соответствующую маску: возможно, имя совпадает или отличается суффиксом\n", "mask_path = None\n", "if img_filename in label_dict:\n", " mask_path = label_dict[img_filename]\n", "else:\n", " # Попробуем убрать суффикс _0000, если он есть\n", " base_name = img_filename.replace('_0000.nii.gz', '.nii.gz')\n", " if base_name in label_dict:\n", " mask_path = label_dict[base_name]\n", "\n", "if mask_path is None:\n", " print(f\"❌ Маска для {img_filename} не найдена. Будет показано только КТ.\")\n", " has_mask = False\n", "else:\n", " has_mask = True\n", " print(f\"✅ Найдена маска: {mask_path}\")\n", "\n", "print(f\"📂 Тестовое изображение: {img_path}\")\n", "\n", "# --- 5. Трансформации для инференса ---\n", "target_shape = [512, 512, 176] # выровненный размер\n", "infer_transform = Compose([\n", " LoadImage(image_only=True),\n", " EnsureChannelFirst(),\n", " ScaleIntensity(),\n", " SpatialPad(spatial_size=target_shape, method=\"symmetric\"),\n", " ToTensor()\n", "])\n", "\n", "input_tensor = infer_transform(img_path).unsqueeze(0).to(device)\n", "\n", "# Если есть маска, загрузим её (без паддинга)\n", "if has_mask:\n", " true_nii = nib.load(mask_path)\n", " true_mask = true_nii.get_fdata() # [H, W, D]\n", " H, W, D = true_mask.shape\n", "else:\n", " # Если маски нет, используем размеры изображения (после паддинга мы не знаем исходный размер, поэтому загрузим оригинал)\n", " orig_nii = nib.load(img_path)\n", " orig_img = orig_nii.get_fdata()\n", " H, W, D = orig_img.shape\n", "\n", "# --- 6. Выбор слоёв для Grad-CAM ---\n", "# Запрашиваем у пользователя индексы через ввод (в Jupyter можно использовать input)\n", "print(\"\\n🔍 Введите индексы слоёв для анализа через пробел (например: 2 3 4):\")\n", "user_input = input(\"Индексы: \").strip()\n", "if user_input:\n", " layer_indices = [int(x) for x in user_input.split()]\n", "else:\n", " # Если ничего не введено, используем последний слой\n", " layer_indices = [len(model.model) - 1]\n", " print(f\" Используется последний слой: {layer_indices[0]}\")\n", "\n", "# Проверяем, что индексы в допустимом диапазоне\n", "max_idx = len(model.model) - 1\n", "layer_indices = [idx for idx in layer_indices if 0 <= idx <= max_idx]\n", "if not layer_indices:\n", " raise ValueError(f\"Нет допустимых индексов. Доступны 0..{max_idx}\")\n", "\n", "print(f\"\\n🔍 Выбранные слои: {layer_indices}\")\n", "\n", "cam_maps = {}\n", "for idx in layer_indices:\n", " target_layer = model.model[idx]\n", " grad_cam = GradCAM(model, target_layer)\n", " cam = grad_cam.generate_cam(input_tensor, target_class=1)\n", " cam_maps[idx] = cam\n", " print(f\" Слой {idx} ({target_layer.__class__.__name__}) – карта получена\")\n", "\n", "# --- 7. Визуализация ---\n", "# Выбираем несколько равномерно распределённых срезов по оси Z\n", "num_slices = 5\n", "slices = np.linspace(0, D-1, num_slices, dtype=int)\n", "\n", "# Если есть маска, визуализируем её контуром, иначе только КТ + Grad-CAM\n", "plt.figure(figsize=(4*num_slices, 4*len(layer_indices)))\n", "\n", "for i, layer_idx in enumerate(layer_indices):\n", " cam_full = cam_maps[layer_idx] # [H_pad, W_pad, D_pad]\n", " for j, slice_idx in enumerate(slices):\n", " # Обрезаем карту до исходного размера\n", " cam_slice = cam_full[:H, :W, slice_idx]\n", " \n", " # Загружаем срез КТ (можно из оригинального файла, но проще из уже загруженного)\n", " # Для КТ используем orig_img (если нет маски, мы его уже загрузили)\n", " if has_mask:\n", " img_slice = true_nii.dataobj[:, :, slice_idx].T # может потребоваться транспонирование\n", " else:\n", " img_slice = orig_img[:, :, slice_idx].T\n", " \n", " # Позиция в подграфике\n", " plt.subplot(len(layer_indices), num_slices, i*num_slices + j + 1)\n", " \n", " # Показываем КТ подложкой\n", " plt.imshow(img_slice, cmap='gray')\n", " # Накладываем Grad-CAM\n", " plt.imshow(cam_slice, cmap='jet', alpha=0.4, interpolation='bilinear')\n", " \n", " # Если есть маска, добавляем её контур\n", " if has_mask:\n", " mask_slice = true_mask[:, :, slice_idx]\n", " plt.contour(mask_slice, colors='white', linewidths=0.5, alpha=0.8)\n", " \n", " plt.title(f'Слой {layer_idx}, срез {slice_idx}')\n", " plt.axis('off')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "DataSphere Kernel", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }