【skull-stripping入門】脳MRIから頭蓋骨を除去！Pythonで実装する手法を解説

# NFBS データセットの取得と解凍（Google Colaboratory）
!wget https://fcp-indi.s3.amazonaws.com/data/Projects/RocklandSample/NFBS_Dataset.tar.gz
!tar -zxvf /content/NFBS_Dataset.tar.gz

import os
import glob
import nibabel as nib
import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

print("全てのライブラリのインポートが完了しました！")

# NFBS データセットのルート
data_dir = "/content/NFBS_Dataset"

# 画像とマスクを格納するリスト
voxels = []  # MRI画像
labels = []  # マスク画像

# 最大50個のデータのみ読み込み
max_samples = 50
sample_count = 0

# 各被験者フォルダをループ
subject_dirs = sorted(glob.glob(os.path.join(data_dir, "A*")))
pbar = tqdm(subject_dirs, total=min(len(subject_dirs), max_samples), desc="読み込み中")
for subject_dir in pbar:
    if sample_count >= max_samples:
        break

    # MRI とマスクのファイルパスを取得
    mri_files  = glob.glob(os.path.join(subject_dir, "*_T1w.nii.gz"))
    mask_files = glob.glob(os.path.join(subject_dir, "*_T1w_brainmask.nii.gz"))

    if not mri_files or not mask_files:
        print(f"ファイルが見つかりません: {subject_dir}")
        continue

    mri_path  = mri_files[0]
    mask_path = mask_files[0]

    # NIfTI を読み込み → numpy 配列に変換
    mri_img_obj = nib.squeeze_image(nib.as_closest_canonical(nib.load(mri_path)))
    mask_img_obj = nib.squeeze_image(nib.as_closest_canonical(nib.load(mask_path)))

    # データ部分を取得してnumpy配列に変換
    mri_data = mri_img_obj.get_fdata().astype("float32")
    mask_data = mask_img_obj.get_fdata().astype("int16")

    # リストに追加
    voxels.append(mri_data)
    labels.append(mask_data)
    sample_count += 1
    pbar.set_postfix({"読み込み済み": sample_count})

print(f"読み込み完了: {len(voxels)} 件")

# 訓練・検証分割（8:2）
train_voxels, val_voxels, train_labels, val_labels = train_test_split(voxels, labels, test_size=0.2, random_state=42)

print(f"訓練データ: {len(train_voxels)} 件, 検証データ: {len(val_voxels)} 件")

import matplotlib.pyplot as plt
import numpy as np

def visualize_slices_horizontal(voxels, labels, view='axial', n_patients=10):
    """
    患者を横並びで可視化（3行×n列）
    view: 'axial', 'sagittal', 'coronal'
    """
    n_patients = min(n_patients, len(voxels))
    n_rows, n_cols = 3, n_patients

    def fix_orientation(image_slice):
        if view == 'axial':
            return np.rot90(image_slice, k=3)
        elif view == 'sagittal':
            return np.fliplr(np.rot90(image_slice, k=1))
        elif view == 'coronal':
            return np.rot90(image_slice, k=1)
        return image_slice

    slice_idx = []
    for v in voxels[:n_patients]:
        if view == 'axial':
            target_slice = min(165, v.shape[2] - 1)
            slice_idx.append(target_slice)
        elif view == 'sagittal':
            slice_idx.append(v.shape[0] // 2)
        elif view == 'coronal':
            slice_idx.append(v.shape[1] // 2)

    fig, axes = plt.subplots(n_rows, n_cols, figsize=(1.5 * n_cols, 6))
    fig.suptitle(f'{view.capitalize()} View', fontsize=16, y=0.95)

    for col in range(n_patients):
        if view == 'axial':
            mri_slice = fix_orientation(voxels[col][:, :, slice_idx[col]])
            mask_slice = fix_orientation(labels[col][:, :, slice_idx[col]])
        elif view == 'sagittal':
            mri_slice = fix_orientation(voxels[col][slice_idx[col], :, :])
            mask_slice = fix_orientation(labels[col][slice_idx[col], :, :])
        else:
            mri_slice = fix_orientation(voxels[col][:, slice_idx[col], :])
            mask_slice = fix_orientation(labels[col][:, slice_idx[col], :])

        axes[0, col].imshow(mri_slice, cmap='gray')
        axes[0, col].axis('off')
        axes[0, col].set_title(f"Sub {col+1}", fontsize=10)
        axes[1, col].imshow(mask_slice, cmap='gray')
        axes[1, col].axis('off')
        axes[2, col].imshow(mri_slice * (mask_slice > 0), cmap='gray')
        axes[2, col].axis('off')

    plt.subplots_adjust(wspace=0.05, hspace=0.05)
    plt.show()

visualize_slices_horizontal(voxels, labels, view='axial', n_patients=10)

visualize_slices_horizontal(voxels, labels, view='sagittal', n_patients=10)

visualize_slices_horizontal(voxels, labels, view='coronal', n_patients=10)

import torch
import numpy as np
from torch.utils.data import Dataset

# カスタムデータセット（2Dスライス用）
class SkullStrippingDataset(Dataset):
    def __init__(self, voxels, labels):
        self.slices = []
        self.masks = []

        # 全てのボリュームから2Dスライスを抽出（Axial方向 = 3軸目）
        for vol, mask in zip(voxels, labels):
            for z in range(vol.shape[2]):
                slice_2d = vol[:, :, z]
                mask_2d = mask[:, :, z]

                # 空のスライスはスキップ
                if np.sum(slice_2d) > 0:
                    self.slices.append(slice_2d)
                    self.masks.append(mask_2d)

        print(f"総スライス数: {len(self.slices)}")

    def __len__(self):
        return len(self.slices)

    def __getitem__(self, idx):
        slice_2d = self.slices[idx]
        mask_2d = self.masks[idx]

        # Min-Max正規化
        slice_min = slice_2d.min()
        slice_max = slice_2d.max()
        if slice_max > slice_min:
            slice_2d = (slice_2d - slice_min) / (slice_max - slice_min)

        # マスクを0-1にクリップ
        mask_2d = np.clip(mask_2d, 0, 1)

        # テンソルに変換（チャンネル次元を追加）
        slice_2d = torch.FloatTensor(slice_2d).unsqueeze(0)  # (1, H, W)
        mask_2d = torch.FloatTensor(mask_2d).unsqueeze(0)    # (1, H, W)

        return slice_2d, mask_2d

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

# UNetモデル
class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super(UNet, self).__init__()

        # エンコーダ
        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)

        # ボトルネック
        self.bottleneck = self.conv_block(512, 1024)

        # デコーダ
        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.dec4 = self.conv_block(1024, 512)

        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = self.conv_block(512, 256)

        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = self.conv_block(256, 128)

        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = self.conv_block(128, 64)

        # 出力層
        self.out = nn.Conv2d(64, out_channels, 1)

    def conv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        # エンコード
        e1 = self.enc1(x)
        e2 = self.enc2(F.max_pool2d(e1, 2))
        e3 = self.enc3(F.max_pool2d(e2, 2))
        e4 = self.enc4(F.max_pool2d(e3, 2))

        # ボトルネック
        b = self.bottleneck(F.max_pool2d(e4, 2))

        # デコード
        d4 = self.up4(b)
        d4 = torch.cat([d4, e4], dim=1)
        d4 = self.dec4(d4)

        d3 = self.up3(d4)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)

        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)

        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)

        # 出力
        out = self.out(d1)
        return out

# デバイス設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用デバイス: {device}")

# データセットとデータローダーの作成
train_dataset = SkullStrippingDataset(train_voxels, train_labels)
val_dataset = SkullStrippingDataset(val_voxels, val_labels)

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)

print(f"訓練スライス数: {len(train_dataset)}, 検証スライス数: {len(val_dataset)}")

import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt

# モデル、損失関数、オプティマイザーの設定
model = UNet().to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 学習ループ
num_epochs = 10
train_losses = []
val_losses = []

print("学習開始...")

for epoch in range(num_epochs):
    # 訓練フェーズ
    model.train()
    train_loss = 0.0

    for batch_idx, (data, target) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    # 検証フェーズ
    model.eval()
    val_loss = 0.0

    with torch.no_grad():
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            val_loss += criterion(output, target).item()

    # 平均損失を計算
    train_loss = train_loss / len(train_loader)
    val_loss = val_loss / len(val_loader)

    train_losses.append(train_loss)
    val_losses.append(val_loss)

    print(f"Epoch {epoch+1}/{num_epochs}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

print("学習完了!")

# モデルの保存
torch.save(model.state_dict(), '/content/NFBS_Dataset/skull_stripping_model.pth')
print("モデルを保存しました: skull_stripping_model.pth")

# 損失の可視化
plt.figure(figsize=(10, 5))
plt.plot(range(1, num_epochs+1), train_losses, 'b-', label='Training Loss')
plt.plot(range(1, num_epochs+1), val_losses, 'r-', label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Validation Loss')
plt.legend()
plt.grid(True)
plt.show()

# ========== 推論用の関数 ==========
def load_trained_model(model_path, device):
    """学習済みモデルを読み込む"""
    model = UNet().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    print(f"モデル読み込み完了: {model_path}")
    return model

def predict_single_volume(model, volume, device, batch_size=8):
    """
    3D画像からskull-stripping予測を実行
    volume: 3D numpy配列 (H, W, D)
    """
    H, W, D = volume.shape
    predicted_mask = np.zeros_like(volume, dtype=np.float32)

    slices = []
    slice_indices = []

    for z in range(D):
        slice_2d = volume[:, :, z].astype(np.float32)

        # Min-Max正規化
        slice_min = slice_2d.min()
        slice_max = slice_2d.max()
        if slice_max > slice_min:
            slice_2d = (slice_2d - slice_min) / (slice_max - slice_min)

        slice_tensor = torch.FloatTensor(slice_2d).unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
        slices.append(slice_tensor)
        slice_indices.append(z)

    # バッチ処理で推論
    with torch.no_grad():
        for i in range(0, len(slices), batch_size):
            batch_slices = slices[i:i+batch_size]
            batch_indices = slice_indices[i:i+batch_size]
            batch_tensor = torch.cat(batch_slices, dim=0).to(device)

            outputs = model(batch_tensor)
            predictions = torch.sigmoid(outputs)

            for j, z_idx in enumerate(batch_indices):
                predicted_mask[:, :, z_idx] = predictions[j, 0].cpu().numpy()

    return predicted_mask

def apply_threshold(prediction, threshold=0.5):
    """予測結果に閾値を適用してバイナリマスクに変換"""
    return (prediction > threshold).astype(np.uint8)

def visualize_results(original, prediction, binary_mask, slice_idx=None):
    """結果を可視化"""
    if slice_idx is None:
        slice_idx = original.shape[2] // 2

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    axes[0].imshow(original[:, :, slice_idx], cmap='gray')
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    axes[1].imshow(prediction[:, :, slice_idx], cmap='jet', vmin=0, vmax=1)
    axes[1].set_title('Prediction Probability')
    axes[1].axis('off')

    axes[2].imshow(original[:, :, slice_idx], cmap='gray', alpha=0.7)
    axes[2].imshow(binary_mask[:, :, slice_idx], cmap='jet', alpha=0.3)
    axes[2].set_title('Binary Mask Overlay')
    axes[2].axis('off')

    plt.tight_layout()
    plt.show()

def save_prediction(prediction, reference_nifti_path, output_path):
    """予測結果をNIfTIファイルとして保存"""
    ref_img = nib.squeeze_image(nib.as_closest_canonical(nib.load(reference_nifti_path)))
    pred_img = nib.Nifti1Image(prediction, ref_img.affine, ref_img.header)
    nib.save(pred_img, output_path)
    print(f"予測結果を保存しました: {output_path}")

# ========== 実際の使用例 ==========

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

# 1. 学習済みモデルの読み込み
model_path = '/content/drive/MyDrive/em/skull_stripping_model.pth'
model = load_trained_model(model_path, device)

# 2. テスト用画像の読み込み
test_image_path = "/content/drive/MyDrive/em/NFBS_Dataset/A00062288/sub-A00062288_ses-NFB3_T1w.nii.gz"
test_img = nib.squeeze_image(nib.as_closest_canonical(nib.load(test_image_path)))
test_volume = test_img.get_fdata().astype("float32")

print(f"テスト画像サイズ: {test_volume.shape}")

# 3. 推論実行
print("推論実行中...")
predicted_mask = predict_single_volume(model, test_volume, device, batch_size=4)

# 4. バイナリマスクに変換
binary_mask = apply_threshold(predicted_mask, threshold=0.5)

# 5. 結果の可視化
print("結果を可視化中...")
visualize_results(test_volume, predicted_mask, binary_mask)

# 6. 結果の保存
output_path = '/content/drive/MyDrive/em/predicted_brainmask.nii.gz'
save_prediction(binary_mask.astype(np.float32), test_image_path, output_path)

print("推論完了！")

# ========== 複数スライスの可視化（オプション） ==========
def visualize_multiple_slices(original, prediction, binary_mask, num_slices=6):
    """複数のスライスを表示"""
    D = original.shape[2]
    slice_indices = np.linspace(D//4, 3*D//4, num_slices).astype(int)

    fig, axes = plt.subplots(3, num_slices, figsize=(20, 12))

    for i, slice_idx in enumerate(slice_indices):
        # 元画像
        axes[0, i].imshow(original[:, :, slice_idx], cmap='gray')
        axes[0, i].set_title(f'Original (z={slice_idx})')
        axes[0, i].axis('off')

        # 予測確率
        axes[1, i].imshow(prediction[:, :, slice_idx], cmap='jet', vmin=0, vmax=1)
        axes[1, i].set_title(f'Prediction (z={slice_idx})')
        axes[1, i].axis('off')

        # オーバーレイ
        axes[2, i].imshow(original[:, :, slice_idx], cmap='gray', alpha=0.7)
        axes[2, i].imshow(binary_mask[:, :, slice_idx], cmap='jet', alpha=0.3)
        axes[2, i].set_title(f'Overlay (z={slice_idx})')
        axes[2, i].axis('off')

    plt.tight_layout()
    plt.show()

# 複数スライス表示
visualize_multiple_slices(test_volume, predicted_mask, binary_mask)

# ========== 単一スライスでの4種類表示 ==========
def visualize_results_four_types(original, prediction, binary_mask, slice_idx=None):
    """単一スライスで4種類の結果を表示"""
    if slice_idx is None:
        slice_idx = original.shape[2] // 2  # 中央のスライス

    original_slice = original[:, :, slice_idx]
    prediction_slice = prediction[:, :, slice_idx]
    mask_slice = binary_mask[:, :, slice_idx]
    skull_stripped = original_slice * (mask_slice > 0)

    fig, axes = plt.subplots(1, 4, figsize=(20, 5))

    # 1. 元画像
    axes[0].imshow(original_slice, cmap='gray')
    axes[0].set_title('Original Image', fontsize=12)
    axes[0].axis('off')

    # 2. 予測確率
    axes[1].imshow(prediction_slice, cmap='hot', vmin=0, vmax=1)
    axes[1].set_title('Prediction Probability', fontsize=12)
    axes[1].axis('off')

    # 3. バイナリマスク
    axes[2].imshow(mask_slice, cmap='gray')
    axes[2].set_title('Binary Mask', fontsize=12)
    axes[2].axis('off')

    # 4. Skull-stripped
    axes[3].imshow(skull_stripped, cmap='gray')
    axes[3].set_title('Skull-stripped Result', fontsize=12)
    axes[3].axis('off')

    plt.tight_layout()
    plt.show()

# ========== 使用例 ==========

# 5. 結果の可視化（4種類表示）
print("結果を可視化中...")
visualize_results_four_types(test_volume, predicted_mask, binary_mask)