ai-med-edu/seminars/Тема 1.2/Т2.ipynb

In [ ]:

%pip install nibabel monai

Cell:

[Cell type raw - unsupported, skipped]

In [ ]:

# Task09_Spleen (рекомендую для начала) - 1.5 ГБ
!wget -O Task09_Spleen.tar 'https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar'

In [ ]:

import tarfile
import os
import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np

# Распаковка
with tarfile.open('Task09_Spleen.tar', 'r') as tar:
    tar.extractall()

# Посмотрим содержимое
base_dir = 'Task09_Spleen'
print("Содержимое:", os.listdir(base_dir))
# Ожидается: imagesTr (тренировочные изображения), labelsTr (маски), imagesTs (тестовые)

In [ ]:

import os

base_dir = 'Task09_Spleen'
print("Папки в датасете:", os.listdir(base_dir))
# Ожидаем: ['imagesTr', 'labelsTr', 'imagesTs', 'dataset.json']

In [ ]:

train_images = sorted(os.listdir(os.path.join(base_dir, 'imagesTr')))
train_labels = sorted(os.listdir(os.path.join(base_dir, 'labelsTr')))

print(f"Тренировочных изображений: {len(train_images)}")
print(f"Масок: {len(train_labels)}")
print(f"Пример имени файла изображения: {train_images[0]}")
print(f"Пример имени файла маски: {train_labels[0]}")

In [ ]:

import os
import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np

base_dir = 'Task09_Spleen'

# Получаем списки файлов, исключая служебные (начинающиеся с ._)
train_images = sorted([
    f for f in os.listdir(os.path.join(base_dir, 'imagesTr')) 
    if not f.startswith('._') and f.endswith('.nii.gz')
])
train_labels = sorted([
    f for f in os.listdir(os.path.join(base_dir, 'labelsTr')) 
    if not f.startswith('._') and f.endswith('.nii.gz')
])

print(f"Тренировочных изображений (реальных): {len(train_images)}")
print(f"Масок (реальных): {len(train_labels)}")
print(f"Пример имени файла изображения: {train_images[0]}")
print(f"Пример имени файла маски: {train_labels[0]}")

# Проверим, совпадают ли имена (в MSD они могут немного отличаться, например, у изображений может быть суффикс '_0000')
# Для простоты возьмём первый файл из изображений и попробуем найти соответствующую маску
# Обычно в Task09_Spleen имена совпадают (без суффиксов)

In [ ]:

# Выбираем первый реальный файл изображения и соответствующую маску
# Предполагаем, что имена совпадают
img_name = train_images[0]
# Ищем маску с таким же именем (в Task09_Spleen они одинаковые)
mask_name = img_name  # если имена совпадают

img_path = os.path.join(base_dir, 'imagesTr', img_name)
mask_path = os.path.join(base_dir, 'labelsTr', mask_name)

# Загрузка
img_nii = nib.load(img_path)
mask_nii = nib.load(mask_path)

img_data = img_nii.get_fdata()
mask_data = mask_nii.get_fdata()

print(f"Размер изображения: {img_data.shape}")
print(f"Уникальные значения в маске: {np.unique(mask_data)}")  # должно быть [0., 1.] (фон и селезёнка)

# Покажем несколько срезов
slices = np.linspace(0, img_data.shape[2]-1, 5, dtype=int)

plt.figure(figsize=(15, 10))
for i, idx in enumerate(slices):
    plt.subplot(2, 5, i+1)
    plt.imshow(img_data[:, :, idx], cmap='gray')
    plt.title(f'Срез {idx}')
    plt.axis('off')

    plt.subplot(2, 5, i+6)
    plt.imshow(mask_data[:, :, idx], cmap='hot', alpha=0.7)
    plt.imshow(img_data[:, :, idx], cmap='gray', alpha=0.3)
    plt.title(f'Маска срез {idx}')
    plt.axis('off')
plt.tight_layout()
plt.show()

In [ ]:

import os
import torch
import nibabel as nib
import numpy as np
from monai.networks.nets import UNet
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.inferers import sliding_window_inference
from monai.transforms import (
    Compose, LoadImaged, EnsureChannelFirstd, ScaleIntensityd,
    RandRotate90d, RandFlipd, RandShiftIntensityd, ToTensord,
    SpatialPadd
)
from monai.data import CacheDataset, DataLoader

# --- 1. Подготовка данных ---
base_dir = 'Task09_Spleen'
images_dir = os.path.join(base_dir, 'imagesTr')
labels_dir = os.path.join(base_dir, 'labelsTr')

train_images = sorted([os.path.join(images_dir, f) 
                       for f in os.listdir(images_dir) 
                       if f.endswith('.nii.gz') and not f.startswith('._')])
train_labels = sorted([os.path.join(labels_dir, f) 
                       for f in os.listdir(labels_dir) 
                       if f.endswith('.nii.gz') and not f.startswith('._')])

assert len(train_images) == len(train_labels), "Количество изображений и масок не совпадает"
print(f"Найдено пар изображение-маска: {len(train_images)}")

# --- 2. Определение максимального размера и округление до кратного 16 ---
print("Определяем максимальный размер по всем осям...")
max_shape = [0, 0, 0]
for img_path in train_images:
    img = nib.load(img_path)
    shape = img.shape
    max_shape = [max(max_shape[i], shape[i]) for i in range(3)]
print(f"Максимальный размер (H, W, D): {max_shape}")

# Округляем до ближайшего числа, кратного 16
aligned_shape = [((dim + 15) // 16) * 16 for dim in max_shape]
print(f"Размер после выравнивания (кратно 16): {aligned_shape}")

# --- 3. Создание словарей данных и разбиение на train/val ---
data_dicts = [{"image": img, "label": lbl} for img, lbl in zip(train_images, train_labels)]
split_idx = int(0.8 * len(data_dicts))
train_dicts = data_dicts[:split_idx]
val_dicts = data_dicts[split_idx:]
print(f"Тренировочных: {len(train_dicts)}, Валидационных: {len(val_dicts)}")

# --- 4. Трансформы (используем aligned_shape) ---
train_transforms = Compose([
    LoadImaged(keys=["image", "label"]),
    EnsureChannelFirstd(keys=["image", "label"]),
    ScaleIntensityd(keys=["image"]),
    SpatialPadd(keys=["image", "label"], spatial_size=aligned_shape, method="symmetric"),
    RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=(0, 1)),
    RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
    RandShiftIntensityd(keys=["image"], offsets=0.1, prob=0.5),
    ToTensord(keys=["image", "label"])
])

val_transforms = Compose([
    LoadImaged(keys=["image", "label"]),
    EnsureChannelFirstd(keys=["image", "label"]),
    ScaleIntensityd(keys=["image"]),
    SpatialPadd(keys=["image", "label"], spatial_size=aligned_shape, method="symmetric"),
    ToTensord(keys=["image", "label"])
])

# --- 5. Датасеты и загрузчики ---
train_ds = CacheDataset(data=train_dicts, transform=train_transforms, cache_rate=0.5)
val_ds = CacheDataset(data=val_dicts, transform=val_transforms, cache_rate=0.5)

batch_size = 2
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2)

# --- 6. Модель, функция потерь, оптимизатор ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Используется устройство: {device}")

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)

loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean")

roi_size = (96, 96, 96)
sw_batch_size = 4

# --- 7. Цикл обучения ---
max_epochs = 50
val_interval = 2
best_metric = -1
best_metric_epoch = -1

for epoch in range(max_epochs):
    print(f"Эпоха {epoch+1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        step += 1
    
    epoch_loss /= step
    print(f"  Средняя потеря: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            for val_data in val_loader:
                val_inputs, val_labels = val_data["image"].to(device), val_data["label"].to(device)
                val_pred = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, model, overlap=0.5
                )
                dice_metric(y_pred=val_pred, y=val_labels)
            metric = dice_metric.aggregate().item()
            dice_metric.reset()
            print(f"  Dice на валидации: {metric:.4f}")
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), "best_spleen_model.pth")
                print("  ✅ Модель сохранена")

print(f"🎉 Лучший Dice: {best_metric:.4f} на эпохе {best_metric_epoch}")

In [ ]:

import torch
from monai.networks.nets import UNet

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Создаём архитектуру (те же параметры, что и при обучении)
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)

# Загружаем веса
checkpoint = torch.load("best_spleen_model.pth", map_location=device)
model.load_state_dict(checkpoint)
model.eval()
print("✅ Модель загружена и готова к инференсу.")

In [ ]:

from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor
import matplotlib.pyplot as plt
import numpy as np

# Параметры (те же, что использовались при обучении)
target_shape = [512, 512, 176]  # выровненный размер

# Трансформы для инференса
infer_transform = Compose([
    LoadImage(image_only=True),
    EnsureChannelFirst(),
    ScaleIntensity(),
    SpatialPad(spatial_size=target_shape, method="symmetric"),
    ToTensor()
])

# Выберите тестовый файл (например, первый из imagesTs)
import os
test_dir = 'Task09_Spleen/imagesTs'
if os.path.exists(test_dir):
    test_files = [f for f in os.listdir(test_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
    if test_files:
        test_path = os.path.join(test_dir, test_files[0])
    else:
        # Если нет тестовых, возьмём валидационный из тренировочных (например, 35-й)
        train_dir = 'Task09_Spleen/imagesTr'
        train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
        test_path = os.path.join(train_dir, train_files[-5])  # последний файл
else:
    train_dir = 'Task09_Spleen/imagesTr'
    train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
    test_path = os.path.join(train_dir, train_files[-5])

print(f"Тестовый файл: {test_path}")

# Применяем трансформации и добавляем batch dimension
input_tensor = infer_transform(test_path).unsqueeze(0).to(device)  # [1,1,H,W,D]

In [ ]:

from monai.inferers import sliding_window_inference

roi_size = (96, 96, 96)   # размер окна (можно подобрать)
sw_batch_size = 4

with torch.no_grad():
    # Прямой проход (если изображение целиком влезает, можно использовать model(input_tensor))
    # Но для 3D лучше sliding window:
    pred_logits = sliding_window_inference(
        input_tensor, roi_size, sw_batch_size, model, overlap=0.5
    )
    pred_mask = torch.argmax(pred_logits, dim=1).cpu().numpy()[0]  # [H,W,D]

print("Предсказание выполнено.")
print(f"Уникальные значения в маске: {np.unique(pred_mask)}")  # должно быть [0,1]

In [ ]:

# Загружаем оригинальное изображение (без трансформаций)
import nibabel as nib
orig_img_nii = nib.load(test_path)
orig_img = orig_img_nii.get_fdata()  # [H,W,D] исходный размер

# Размеры
H, W, D = orig_img.shape
print(f"Исходный размер: {orig_img.shape}")

# Обрежем предсказанную маску до исходного размера (если нужно)
pred_mask_cropped = pred_mask[:H, :W, :D]

# Выберем несколько срезов по оси Z
slices = np.linspace(0, D-1, 5, dtype=int)

plt.figure(figsize=(15, 10))
for i, idx in enumerate(slices):
    plt.subplot(2, 5, i+1)
    plt.imshow(orig_img[:, :, idx], cmap='gray')
    plt.title(f'Срез {idx}')
    plt.axis('off')

    plt.subplot(2, 5, i+6)
    plt.imshow(pred_mask_cropped[:, :, idx], cmap='hot', alpha=0.7)
    plt.imshow(orig_img[:, :, idx], cmap='gray', alpha=0.3)
    plt.title(f'Маска срез {idx}')
    plt.axis('off')
plt.tight_layout()
plt.show()

In [ ]:

import torch
import numpy as np
from monai.networks.nets import UNet
from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric, HausdorffDistanceMetric
import os
import nibabel as nib

# 1. Загружаем модель (как вы уже сделали)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
checkpoint = torch.load("best_spleen_model.pth", map_location=device)
model.load_state_dict(checkpoint)
model.eval()

# 2. Параметры (те же, что при обучении)
target_shape = [512, 512, 176]  # выровненный размер
roi_size = (96, 96, 96)
sw_batch_size = 4

# Трансформы для инференса
infer_transform = Compose([
    LoadImage(image_only=True),
    EnsureChannelFirst(),
    ScaleIntensity(),
    SpatialPad(spatial_size=target_shape, method="symmetric"),
    ToTensor()
])

# 3. Собираем все валидационные пары (изображение + маска)
base_dir = 'Task09_Spleen'
val_images = sorted([os.path.join(base_dir, 'imagesTr', f) 
                     for f in os.listdir(os.path.join(base_dir, 'imagesTr')) 
                     if f.endswith('.nii.gz') and not f.startswith('._')])
val_labels = sorted([os.path.join(base_dir, 'labelsTr', f) 
                     for f in os.listdir(os.path.join(base_dir, 'labelsTr')) 
                     if f.endswith('.nii.gz') and not f.startswith('._')])

# Предположим, что последние 9 файлов были валидационными (как в вашем разбиении)
val_images = val_images[-9:]
val_labels = val_labels[-9:]

print(f"Валидационных образцов: {len(val_images)}")

# 4. Метрики
dice_metric = DiceMetric(include_background=False, reduction="mean")
hausdorff_metric = HausdorffDistanceMetric(include_background=False, percentile=95)

# 5. Цикл по всем валидационным образцам
all_dice = []
all_hd95 = []

with torch.no_grad():
    for img_path, mask_path in zip(val_images, val_labels):
        # Подготовка изображения
        input_tensor = infer_transform(img_path).unsqueeze(0).to(device)
        
        # Предсказание
        pred_logits = sliding_window_inference(
            input_tensor, roi_size, sw_batch_size, model, overlap=0.5
        )
        pred_mask = torch.argmax(pred_logits, dim=1).cpu().numpy()[0]  # [H,W,D]
        
        # Загрузка истинной маски (без паддинга)
        true_mask_nii = nib.load(mask_path)
        true_mask = true_mask_nii.get_fdata()
        
        # Обрезаем предсказание до исходного размера
        H, W, D = true_mask.shape
        pred_mask_cropped = pred_mask[:H, :W, :D]
        
        # Преобразуем в тензоры для MONAI метрик
        pred_tensor = torch.from_numpy(pred_mask_cropped).unsqueeze(0).unsqueeze(0).long()  # [1,1,H,W,D]
        true_tensor = torch.from_numpy(true_mask).unsqueeze(0).unsqueeze(0).long()
        
        # Вычисляем Dice
        dice = dice_metric(pred_tensor, true_tensor)
        all_dice.append(dice.item())
        
        # Вычисляем Hausdorff Distance 95%
        hd = hausdorff_metric(pred_tensor, true_tensor)
        all_hd95.append(hd.item())
        
        print(f"Файл: {os.path.basename(img_path)} -> Dice: {dice.item():.4f}, HD95: {hd.item():.2f} мм")

# 6. Итоговая статистика
print("\n" + "="*40)
print(f"Средний Dice по валидации: {np.mean(all_dice):.4f} ± {np.std(all_dice):.4f}")
print(f"Средний HD95 по валидации: {np.mean(all_hd95):.2f} ± {np.std(all_hd95):.2f} мм")
print("="*40)

In [ ]:

import torch
import numpy as np
import matplotlib.pyplot as plt
from monai.networks.nets import UNet
from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor
from monai.inferers import sliding_window_inference
import nibabel as nib
import os

# --- 1. Загрузка модели ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
checkpoint = torch.load("best_spleen_model.pth", map_location=device)
model.load_state_dict(checkpoint)
model.eval()
print("✅ Модель загружена")

# --- 2. Вывод структуры model.model для выбора целевого слоя ---
print("\nСтруктура model.model:")
for idx, module in enumerate(model.model):
    print(f"[{idx}]: {module.__class__.__name__}")

# --- 3. Выбор целевого слоя ---
# Обычно для Grad‑CAM берут последний сверточный слой энкодера.
# В U‑Net от MONAI энкодер состоит из блоков `Down`.
# Выберите индекс блока, который вы хотите использовать (например, последний Down перед bottleneck).
# Попробуйте индекс 2 или 3, если они есть.
# В вашем случае, возможно, подойдёт индекс 2.
# Если не уверены, начните с последнего Down блока.
target_layer_idx = 2  # Попробуйте изменить на 2 или 3, в зависимости от вывода выше
target_layer = model.model[target_layer_idx]
print(f"\nВыбран целевой слой: model.model[{target_layer_idx}] -> {target_layer.__class__.__name__}")

# --- 4. Класс Grad-CAM ---
class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        target_layer.register_forward_hook(self.save_activation)
        target_layer.register_full_backward_hook(self.save_gradient)  # обратите внимание на register_full_backward_hook
    
    def save_activation(self, module, input, output):
        self.activations = output.detach()
    
    def save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()
    
    def generate_cam(self, input_tensor, target_class=1):
        self.model.zero_grad()
        output = self.model(input_tensor)
        
        target = output[:, target_class].sum()
        target.backward()
        
        weights = torch.mean(self.gradients, dim=(2, 3, 4), keepdim=True)
        cam = torch.sum(weights * self.activations, dim=1, keepdim=True)
        cam = torch.relu(cam)
        cam = cam / (cam.max() + 1e-8)
        return cam.squeeze().cpu().numpy()

# --- 5. Подготовка тестового изображения ---
test_dir = 'Task09_Spleen/imagesTs'
if os.path.exists(test_dir):
    test_files = [f for f in os.listdir(test_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
    if test_files:
        test_path = os.path.join(test_dir, test_files[0])
    else:
        train_dir = 'Task09_Spleen/imagesTr'
        train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
        test_path = os.path.join(train_dir, train_files[-5])
else:
    train_dir = 'Task09_Spleen/imagesTr'
    train_files = [f for f in os.listdir(train_dir) if f.endswith('.nii.gz') and not f.startswith('._')]
    test_path = os.path.join(train_dir, train_files[-5])

print(f"\nТестовый файл: {test_path}")

target_shape = [512, 512, 176]
infer_transform = Compose([
    LoadImage(image_only=True),
    EnsureChannelFirst(),
    ScaleIntensity(),
    SpatialPad(spatial_size=target_shape, method="symmetric"),
    ToTensor()
])

input_tensor = infer_transform(test_path).unsqueeze(0).to(device)

# --- 6. Получение Grad-CAM ---
grad_cam = GradCAM(model, target_layer)
cam_3d = grad_cam.generate_cam(input_tensor, target_class=1)

# --- 7. Визуализация ---
orig_img = nib.load(test_path).get_fdata()
H, W, D = orig_img.shape

slices = np.linspace(0, D-1, 5, dtype=int)

plt.figure(figsize=(15, 10))
for i, idx in enumerate(slices):
    cam_slice = cam_3d[:H, :W, idx]
    
    plt.subplot(2, 5, i+1)
    plt.imshow(orig_img[:, :, idx], cmap='gray')
    plt.title(f'Срез {idx}')
    plt.axis('off')
    
    plt.subplot(2, 5, i+6)
    plt.imshow(orig_img[:, :, idx], cmap='gray', alpha=0.5)
    plt.imshow(cam_slice, cmap='jet', alpha=0.5, interpolation='bilinear')
    plt.title(f'Grad-CAM срез {idx}')
    plt.axis('off')
plt.tight_layout()
plt.show()

In [ ]:

import torch
import numpy as np
import matplotlib.pyplot as plt
from monai.networks.nets import UNet
from monai.transforms import Compose, LoadImage, EnsureChannelFirst, ScaleIntensity, SpatialPad, ToTensor
import nibabel as nib
import os

# --- 1. Загрузка модели ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
checkpoint = torch.load("best_spleen_model.pth", map_location=device)
model.load_state_dict(checkpoint)
model.eval()
print("✅ Модель загружена")

# --- 2. Вывод структуры model.model ---
print("\n📋 Доступные слои в model.model:")
for idx, module in enumerate(model.model):
    print(f"  [{idx}]: {module.__class__.__name__}")

# --- 3. Класс Grad-CAM ---
class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        target_layer.register_forward_hook(self.save_activation)
        target_layer.register_full_backward_hook(self.save_gradient)
    
    def save_activation(self, module, input, output):
        self.activations = output.detach()
    
    def save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()
    
    def generate_cam(self, input_tensor, target_class=1):
        self.model.zero_grad()
        output = self.model(input_tensor)
        
        target = output[:, target_class].sum()
        target.backward()
        
        weights = torch.mean(self.gradients, dim=(2, 3, 4), keepdim=True)
        cam = torch.sum(weights * self.activations, dim=1, keepdim=True)
        cam = torch.relu(cam)
        cam = cam / (cam.max() + 1e-8)
        return cam.squeeze().cpu().numpy()  # [H, W, D]

# --- 4. Подготовка тестовых данных ---
base_dir = 'Task09_Spleen'
img_dir = os.path.join(base_dir, 'imagesTr')  # используем тренировочные, если нет test
label_dir = os.path.join(base_dir, 'labelsTr')

# Получаем списки файлов
img_files = sorted([f for f in os.listdir(img_dir) if f.endswith('.nii.gz') and not f.startswith('._')])
label_files = sorted([f for f in os.listdir(label_dir) if f.endswith('.nii.gz') and not f.startswith('._')])

print(f"\n📂 Найдено изображений: {len(img_files)}, масок: {len(label_files)}")

# Создаём словарь соответствий: имя маски (без пути) -> полный путь
label_dict = {os.path.basename(f): os.path.join(label_dir, f) for f in label_files}

# Выбираем файл изображения (например, 10-й)
idx_file = 10  # можно изменить
img_filename = img_files[idx_file]
img_path = os.path.join(img_dir, img_filename)

# Ищем соответствующую маску: возможно, имя совпадает или отличается суффиксом
mask_path = None
if img_filename in label_dict:
    mask_path = label_dict[img_filename]
else:
    # Попробуем убрать суффикс _0000, если он есть
    base_name = img_filename.replace('_0000.nii.gz', '.nii.gz')
    if base_name in label_dict:
        mask_path = label_dict[base_name]

if mask_path is None:
    print(f"❌ Маска для {img_filename} не найдена. Будет показано только КТ.")
    has_mask = False
else:
    has_mask = True
    print(f"✅ Найдена маска: {mask_path}")

print(f"📂 Тестовое изображение: {img_path}")

# --- 5. Трансформации для инференса ---
target_shape = [512, 512, 176]  # выровненный размер
infer_transform = Compose([
    LoadImage(image_only=True),
    EnsureChannelFirst(),
    ScaleIntensity(),
    SpatialPad(spatial_size=target_shape, method="symmetric"),
    ToTensor()
])

input_tensor = infer_transform(img_path).unsqueeze(0).to(device)

# Если есть маска, загрузим её (без паддинга)
if has_mask:
    true_nii = nib.load(mask_path)
    true_mask = true_nii.get_fdata()  # [H, W, D]
    H, W, D = true_mask.shape
else:
    # Если маски нет, используем размеры изображения (после паддинга мы не знаем исходный размер, поэтому загрузим оригинал)
    orig_nii = nib.load(img_path)
    orig_img = orig_nii.get_fdata()
    H, W, D = orig_img.shape

# --- 6. Выбор слоёв для Grad-CAM ---
# Запрашиваем у пользователя индексы через ввод (в Jupyter можно использовать input)
print("\n🔍 Введите индексы слоёв для анализа через пробел (например: 2 3 4):")
user_input = input("Индексы: ").strip()
if user_input:
    layer_indices = [int(x) for x in user_input.split()]
else:
    # Если ничего не введено, используем последний слой
    layer_indices = [len(model.model) - 1]
    print(f"   Используется последний слой: {layer_indices[0]}")

# Проверяем, что индексы в допустимом диапазоне
max_idx = len(model.model) - 1
layer_indices = [idx for idx in layer_indices if 0 <= idx <= max_idx]
if not layer_indices:
    raise ValueError(f"Нет допустимых индексов. Доступны 0..{max_idx}")

print(f"\n🔍 Выбранные слои: {layer_indices}")

cam_maps = {}
for idx in layer_indices:
    target_layer = model.model[idx]
    grad_cam = GradCAM(model, target_layer)
    cam = grad_cam.generate_cam(input_tensor, target_class=1)
    cam_maps[idx] = cam
    print(f"  Слой {idx} ({target_layer.__class__.__name__}) – карта получена")

# --- 7. Визуализация ---
# Выбираем несколько равномерно распределённых срезов по оси Z
num_slices = 5
slices = np.linspace(0, D-1, num_slices, dtype=int)

# Если есть маска, визуализируем её контуром, иначе только КТ + Grad-CAM
plt.figure(figsize=(4*num_slices, 4*len(layer_indices)))

for i, layer_idx in enumerate(layer_indices):
    cam_full = cam_maps[layer_idx]  # [H_pad, W_pad, D_pad]
    for j, slice_idx in enumerate(slices):
        # Обрезаем карту до исходного размера
        cam_slice = cam_full[:H, :W, slice_idx]
        
        # Загружаем срез КТ (можно из оригинального файла, но проще из уже загруженного)
        # Для КТ используем orig_img (если нет маски, мы его уже загрузили)
        if has_mask:
            img_slice = true_nii.dataobj[:, :, slice_idx].T  # может потребоваться транспонирование
        else:
            img_slice = orig_img[:, :, slice_idx].T
        
        # Позиция в подграфике
        plt.subplot(len(layer_indices), num_slices, i*num_slices + j + 1)
        
        # Показываем КТ подложкой
        plt.imshow(img_slice, cmap='gray')
        # Накладываем Grad-CAM
        plt.imshow(cam_slice, cmap='jet', alpha=0.4, interpolation='bilinear')
        
        # Если есть маска, добавляем её контур
        if has_mask:
            mask_slice = true_mask[:, :, slice_idx]
            plt.contour(mask_slice, colors='white', linewidths=0.5, alpha=0.8)
        
        plt.title(f'Слой {layer_idx}, срез {slice_idx}')
        plt.axis('off')
plt.tight_layout()
plt.show()

In [ ]: