1# 下載 repository
2git clone https://github.com/shaoyanpan/Synthetic-CT-generation-from-MRI-using-3D-transformer-based-denoising-diffusion-model.git
3cd Synthetic-CT-generation-from-MRI-using-3D-transformer-based-denoising-diffusion-model
4
5# 以 environment.yml 建立 Conda 環境
6conda env create -f environment.yml
7conda activate DL

 1# 建立乾淨環境
 2conda create -n synth-ct python=3.8 -y
 3conda activate synth-ct
 4
 5# 核心套件
 6conda install pytorch=1.9.1 cudatoolkit=11.1 -c pytorch -y
 7pip install monai==0.7.0 nibabel pydicom scipy numpy==1.18.5
 8pip install timm==0.4.12 einops matplotlib h5py
 9pip install simpleitk medpy torchio
10
11# 驗證
12python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
13python -c "import monai; print(f'MONAI {monai.__version__}')"

1data/
2  imagesTr/          # 訓練集
3    patient_1.mat    # dict: {'image': MRI_array, 'label': CT_array}
4    patient_2.mat
5  imagesTs/          # 測試集
6    patient_3.mat
7  imagesVal/         # 驗證集
8    patient_4.mat

1data/
2  imagesTr/          # 訓練 MRI
3    patient_1.nii.gz
4  labelsTr/          # 訓練 CT（ground truth）
5    patient_1.nii.gz
6  imagesTs/          # 測試 MRI
7  labelsTs/          # 測試 CT
8  imagesVal/         # 驗證 MRI
9  labelsVal/         # 驗證 CT

 1import torch
 2from diffusion.Create_diffusion import create_gaussian_diffusion
 3from diffusion.resampler import UniformSampler
 4from network.Diffusion_model_transformer import SwinVITModel
 5
 6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 7
 8# === Step 1: 建立 Gaussian Diffusion ===
 9# 訓練使用 1000 步，推論透過 timestep_respacing 壓縮到 50 步
10diffusion = create_gaussian_diffusion(
11    steps=1000,                    # 總擴散步數
12    learn_sigma=True,              # 學習 variance（輸出 2 通道）
13    sigma_small=False,
14    noise_schedule='linear',       # 線性 beta schedule
15    use_kl=False,                  # 使用 MSE loss（非 KL）
16    predict_xstart=False,          # 預測 noise epsilon（非 x_0）
17    rescale_timesteps=True,
18    rescale_learned_sigmas=True,
19    timestep_respacing=[50],       # 推論時壓縮為 50 步
20)
21
22# 均勻採樣器（訓練時隨機選取 timestep）
23schedule_sampler = UniformSampler(diffusion)
24
25# === Step 2: 建立 3D Swin-Transformer UNet ===
26img_size = (128, 128, 64)  # 3D patch size (D, H, W)
27
28model = SwinVITModel(
29    image_size=img_size,
30    in_channels=2,                  # MRI (condition) + noisy CT
31    model_channels=64,              # 基礎通道數
32    out_channels=2,                 # predicted noise + learned sigma
33    dims=3,                         # 3D convolutions
34    sample_kernel=([2,2,2], [2,2,1], [2,2,1], [2,2,1]),  # downsampling kernels
35    num_res_blocks=[2, 2, 2, 2],    # 每個解析度的 ResBlock 數量
36    attention_resolutions=(32, 16, 8),
37    dropout=0,
38    channel_mult=(1, 2, 3, 4),      # 通道倍率: 64→128→192→256
39    num_classes=None,                # 無 class conditioning
40    use_checkpoint=False,
41    use_fp16=False,
42    num_heads=[4, 4, 8, 16],        # 每層 attention heads
43    window_size=[[4,4,4], [4,4,4], [4,4,2], [4,4,2]],  # 3D window sizes
44    num_head_channels=64,
45    use_scale_shift_norm=True,       # AdaGN (timestep-conditioned normalization)
46    resblock_updown=False,
47).to(device)
48
49print(f"模型參數量: {sum(p.numel() for p in model.parameters()):,}")
50# 預期約 ~50-80M 參數，視 img_size 而定

 1import numpy as np
 2from torch.utils.data import DataLoader, Dataset
 3import scipy.io as sio
 4
 5# === 自定義 Dataset（MATLAB 格式）===
 6class MRI_CT_Dataset(Dataset):
 7    """載入 .mat 檔案，每個檔案包含 'image'（MRI）和 'label'（CT）"""
 8    def __init__(self, data_dir, img_size=(128, 128, 64)):
 9        import glob
10        self.files = sorted(glob.glob(f"{data_dir}/*.mat"))
11        self.img_size = img_size
12    
13    def __len__(self):
14        return len(self.files)
15    
16    def __getitem__(self, idx):
17        data = sio.loadmat(self.files[idx])
18        mri = data['image'].astype(np.float32)   # MRI condition
19        ct = data['label'].astype(np.float32)     # CT target
20        
21        # 正規化到 [-1, 1]
22        mri = (mri - mri.mean()) / (mri.std() + 1e-8)
23        ct = (ct - ct.min()) / (ct.max() - ct.min() + 1e-8) * 2 - 1
24        
25        # 轉為 tensor: (1, D, H, W)
26        mri = torch.tensor(mri).unsqueeze(0)
27        ct = torch.tensor(ct).unsqueeze(0)
28        return mri, ct
29
30# === 訓練迴圈 ===
31dataset = MRI_CT_Dataset("MRI_to_CT_brain_for_dosimetric/imagesTr")
32dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
33
34optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
35
36num_epochs = 100
37for epoch in range(num_epochs):
38    model.train()
39    epoch_loss = 0.0
40    
41    for batch_idx, (mri, ct) in enumerate(dataloader):
42        mri, ct = mri.to(device), ct.to(device)
43        
44        # 從 schedule_sampler 抽取隨機 timestep
45        batch_size = ct.shape[0]
46        t, weights = schedule_sampler.sample(batch_size, device)
47        
48        # 計算 diffusion training loss
49        # target = ct (ground truth CT), condition = mri
50        losses = diffusion.training_losses(
51            model,
52            x_start=ct,          # 目標影像 (CT)
53            t=t,                 # 隨機 timestep
54            model_kwargs={"condition": mri}  # MRI 作為條件
55        )
56        
57        loss = (losses["loss"] * weights).mean()
58        
59        optimizer.zero_grad()
60        loss.backward()
61        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
62        optimizer.step()
63        
64        epoch_loss += loss.item()
65    
66    avg_loss = epoch_loss / len(dataloader)
67    if (epoch + 1) % 10 == 0:
68        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.6f}")
69        # 定期儲存 checkpoint
70        torch.save({
71            'epoch': epoch,
72            'model_state_dict': model.state_dict(),
73            'optimizer_state_dict': optimizer.state_dict(),
74            'loss': avg_loss,
75        }, f"checkpoint_epoch_{epoch+1}.pt")

 1from monai.inferers import SlidingWindowInferer
 2import matplotlib.pyplot as plt
 3
 4# === 載入已訓練的模型 ===
 5checkpoint = torch.load("checkpoint_epoch_100.pt", map_location=device)
 6model.load_state_dict(checkpoint['model_state_dict'])
 7model.eval()
 8
 9# === 設定 SlidingWindowInferer ===
10# 因為訓練時使用固定 patch size，推論時用 sliding window 處理完整體積
11roi_size = (128, 128, 64)   # 與訓練時的 img_size 一致
12sw_batch_size = 4            # 同時處理的 patch 數量
13overlap = 0.5                # patch 間 50% overlap（減少邊界偽影）
14
15inferer = SlidingWindowInferer(
16    roi_size=roi_size,
17    sw_batch_size=sw_batch_size,
18    overlap=overlap,
19    mode='constant'           # overlap 區域的融合策略
20)
21
22# === 定義 diffusion sampling 函數 ===
23def diffusion_sampling(condition, model):
24    """
25    從純雜訊開始，以 MRI condition 引導，迭代去噪生成 Synthetic CT。
26    推論時使用 50 步 (timestep_respacing=[50])，而非完整 1000 步。
27    """
28    with torch.no_grad():
29        sampled = diffusion.p_sample_loop(
30            model,
31            shape=(
32                condition.shape[0], 1,
33                condition.shape[2], condition.shape[3], condition.shape[4]
34            ),
35            condition=condition,
36            clip_denoised=True,   # 將輸出 clip 到 [-1, 1]
37        )
38    return sampled
39
40# === 對測試 MRI 執行推論 ===
41test_dataset = MRI_CT_Dataset("MRI_to_CT_brain_for_dosimetric/imagesTs")
42test_mri, test_ct_real = test_dataset[0]
43test_mri = test_mri.unsqueeze(0).to(device)  # (1, 1, D, H, W)
44
45# SlidingWindowInferer 自動將大體積切成 patches 並合併結果
46synth_ct = inferer(test_mri, diffusion_sampling, model)
47
48# === 視覺化結果 ===
49slice_idx = synth_ct.shape[2] // 2  # 取中間 axial slice
50
51fig, axes = plt.subplots(1, 3, figsize=(15, 5))
52axes[0].imshow(test_mri[0, 0, slice_idx].cpu().numpy(), cmap='gray')
53axes[0].set_title('MRI (Input)')
54axes[1].imshow(synth_ct[0, 0, slice_idx].cpu().numpy(), cmap='gray')
55axes[1].set_title('Synthetic CT (Generated)')
56axes[2].imshow(test_ct_real[0, slice_idx].numpy(), cmap='gray')
57axes[2].set_title('Real CT (Ground Truth)')
58
59for ax in axes:
60    ax.axis('off')
61plt.tight_layout()
62plt.savefig('synth_ct_comparison.png', dpi=150, bbox_inches='tight')
63plt.show()
64
65# === 量化評估 ===
66from skimage.metrics import structural_similarity as ssim
67from skimage.metrics import peak_signal_noise_ratio as psnr
68
69synth_np = synth_ct[0, 0].cpu().numpy()
70real_np = test_ct_real[0].numpy()
71
72mae = np.mean(np.abs(synth_np - real_np))
73psnr_val = psnr(real_np, synth_np, data_range=2.0)  # [-1, 1] range
74ssim_val = ssim(real_np, synth_np, data_range=2.0)
75
76print(f"MAE: {mae:.4f}")
77print(f"PSNR: {psnr_val:.2f} dB")
78print(f"SSIM: {ssim_val:.4f}")

 1Domain 2 生態系定位：
 2
 3[A] 基礎平台層
 4  MONAI / MONAI GenerativeModels ──→ 本專案直接使用 MONAI SlidingWindowInferer
 5  nnUNet                          ──→ 資料組織格式相近
 6
 7[C] 跨模態合成（本專案所在）
 8  medSynthesisV1 (GAN baseline)   ──→ 本專案為其 Diffusion 升級版
 9  medSynthesis (unpaired)         ──→ 互補：paired vs. unpaired 場景
10
11[D] Diffusion 前沿
12  conditional_DDPM (2D baseline)  ──→ 本專案為其 3D + Transformer 擴展
13  DenseDiffusion                  ──→ layout-guided，不同任務但共享 diffusion 基礎
14
15[F] 癌症影像
16  foundation-cancer-image-biomarker ──→ 合成 CT 可做為 downstream biomarker 輸入

 1# 使用 MONAI Transforms 做資料前處理
 2from monai.transforms import (
 3    Compose, LoadImaged, EnsureChannelFirstd,
 4    ScaleIntensityd, CropForegroundd, RandSpatialCropd
 5)
 6
 7train_transforms = Compose([
 8    LoadImaged(keys=["image", "label"]),       # 載入 NIfTI
 9    EnsureChannelFirstd(keys=["image", "label"]),
10    ScaleIntensityd(keys=["image"], minv=-1, maxv=1),
11    ScaleIntensityd(keys=["label"], minv=-1, maxv=1),
12    CropForegroundd(keys=["image", "label"], source_key="image"),
13    RandSpatialCropd(keys=["image", "label"], roi_size=(128, 128, 64)),
14])

 1# 概念性整合：用 MONAI GenerativeModels 的 scheduler 替換手寫 diffusion
 2from generative.networks.schedulers import DDPMScheduler
 3
 4scheduler = DDPMScheduler(
 5    num_train_timesteps=1000,
 6    beta_schedule="linear",
 7    prediction_type="epsilon"
 8)
 9
10# 搭配本專案的 SwinVITModel 作為 denoiser
11# scheduler.add_noise() 取代手寫的 forward process
12# scheduler.step() 取代手寫的 reverse sampling