929 lines
40 KiB
Plaintext
929 lines
40 KiB
Plaintext
{
|
||
"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
|
||
}
|