ResViT 代码复现和讲解

论文题目：ResViT: Residual vision transformers for multi-modal medical image synthesis

论文地址：[2106.16031] ResViT: Residual vision transformers for multi-modal medical image synthesis (arxiv.org)

项目地址：GitHub - icon-lab/ResViT: Official Implementation of ResViT: Residual Vision Transformers for Multi-modal Medical Image Synthesis

一种新的用于多模态医学图像合成的生成对抗方法

复现

1、从github下载源码 

2、数据问题：

本次实验使用的数据不同于原项目，因此给出了新的dataset 方法和预处理方法

BraTS2020 Dataset (Training + Validation)

BraTS2020 Dataset (Training + Validation) (kaggle.com)

下载后解压

3、预处理，

预处理方法使用的MTNet的项目方法，代码讲解可以看我之前的博客

import numpy as np
from matplotlib import pylab as plt
import nibabel as nib
import random
import glob
import os
from PIL import Image
import imageio

def normalize(image, mask=None, percentile_lower=0.2, percentile_upper=99.8):

    if mask is None:
        mask = image != image[0, 0, 0]
    cut_off_lower = np.percentile(image[mask != 0].ravel(), percentile_lower)
    cut_off_upper = np.percentile(image[mask != 0].ravel(), percentile_upper)
    res = np.copy(image)
    res[(res < cut_off_lower) & (mask != 0)] = cut_off_lower
    res[(res > cut_off_upper) & (mask != 0)] = cut_off_upper
    res = res / res.max()  # 0-1

    return res

def visualize(t1_data,t2_data,flair_data,t1ce_data,gt_data):

    plt.figure(figsize=(8, 8))
    plt.subplot(231)
    plt.imshow(t1_data[:, :], cmap='gray')
    plt.title('Image t1')
    plt.subplot(232)
    plt.imshow(t2_data[:, :], cmap='gray')
    plt.title('Image t2')
    plt.subplot(233)
    plt.imshow(flair_data[:, :], cmap='gray')
    plt.title('Image flair')
    plt.subplot(234)
    plt.imshow(t1ce_data[:, :], cmap='gray')
    plt.title('Image t1ce')
    plt.subplot(235)
    plt.imshow(gt_data[:, :])
    plt.title('GT')
    plt.show()

def visualize_to_gif(t1_data, t2_data, t1ce_data, flair_data):
    transversal = []
    coronal = []
    sagittal = []
    slice_num = t1_data.shape[2]
    for i in range(slice_num):
        sagittal_plane = np.concatenate((t1_data[:, :, i], t2_data[:, :, i],
                              t1ce_data[:, :, i],flair_data[:, :, i]),axis=1)
        coronal_plane = np.concatenate((t1_data[i, :, :], t2_data[i, :, :],
                              t1ce_data[i, :, :],flair_data[i, :, :]),axis=1)
        transversal_plane = np.concatenate((t1_data[:, i, :], t2_data[:, i, :],
                              t1ce_data[:, i, :],flair_data[:, i, :]),axis=1)
        transversal.append(transversal_plane)
        coronal.append(coronal_plane)
        sagittal.append(sagittal_plane)
    imageio.mimsave("./transversal_plane.gif", transversal, duration=0.01)
    imageio.mimsave("./coronal_plane.gif", coronal, duration=0.01)
    imageio.mimsave("./sagittal_plane.gif", sagittal, duration=0.01)
    return

if __name__ == '__main__':

    t1_list = sorted(glob.glob(
        '../data/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*/*t1.*'))
    t2_list = sorted(glob.glob(
        '../data/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*/*t2.*'))
    t1ce_list = sorted(glob.glob(
        '../data/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*/*t1ce.*'))
    flair_list = sorted(glob.glob(
        '../data/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*/*flair.*'))
    gt_list = sorted(glob.glob(
        '../data/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/*/*seg.*'))

    data_len = len(gt_list)
    train_len = int(data_len * 0.8)
    test_len = data_len - train_len

    train_path = '../data/train/'
    test_path = '../data/test/'

    os.makedirs(train_path,exist_ok=True)
    os.makedirs(test_path,exist_ok=True)
    for i,(t1_path, t2_path, t1ce_path, flair_path, gt_path) in enumerate(zip(t1_list,t2_list,t1ce_list,flair_list,gt_list)):

        print('preprocessing the',i+1,'th subject')

        t1_img = nib.load(t1_path)  # (240,140,155)
        t2_img = nib.load(t2_path)
        flair_img = nib.load(flair_path)
        t1ce_img = nib.load(t1ce_path)
        gt_img = nib.load(gt_path)

        #to numpy
        t1_data = t1_img.get_fdata()
        t2_data = t2_img.get_fdata()
        flair_data = flair_img.get_fdata()
        t1ce_data = t1ce_img.get_fdata()
        gt_data = gt_img.get_fdata()
        gt_data = gt_data.astype(np.uint8)
        gt_data[gt_data == 4] = 3 #label 3 is missing in BraTS 2020

        t1_data = normalize(t1_data) # normalize to [0,1]
        t2_data = normalize(t2_data)
        t1ce_data = normalize(t1ce_data)
        flair_data = normalize(flair_data)

        tensor = np.stack([t1_data, t2_data, t1ce_data, flair_data, gt_data])  # (4, 240, 240, 155)

        if i < train_len:
            for j in range(60):
                Tensor = tensor[:, 10:210, 25:225, 50 + j]
                np.save(train_path + str(60 * i + j + 1) + '.npy', Tensor)
        else:
            for j in range(60):
                Tensor = tensor[:, 10:210, 25:225, 50 + j]
                np.save(test_path + str(60 * (i - train_len) + j + 1) + '.npy', Tensor)

调整一下文件地址，处理好的如下

每一个npy都是一个 多模态切片 （5，200，200）

分别是t1_data, t2_data, t1ce_data, flair_data, gt_data

4、修改dataset

data/aligned_dataset.py

任务可以根据调整AB 修改

import os.path
import random
# import torchvision.transforms as transforms
import torch
from data.base_dataset import BaseDataset
# from data.image_folder import make_dataset
# from PIL import Image
import numpy as np


class AlignedDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot
        self.pathlist = os.listdir((self.root))
        # self.dir_AB = os.path.join(opt.dataroot, opt.phase)
        # self.AB_paths = sorted(make_dataset(self.dir_AB))
        # assert(opt.resize_or_crop == 'resize_and_crop')

    def __getitem__(self, index):
        casepath = self.pathlist[index]
        data = np.load(os.path.join(self.root,casepath))
        # [t1_data, t2_data, t1ce_data, flair_data, gt_data]

        A = data.take([0,1,3],0)
        B = data[2,:,:]
        B = np.expand_dims(B,0)


        w_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))
        h_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))
        A = A[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]
        B = B[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]


        A = torch.tensor(A).float()
        B = torch.tensor(B).float()

        return {'A': A, 'B': B, 'A_paths': casepath, 'B_paths':casepath}

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

    def name(self):
        return 'AlignedDataset'

5、运行！

nohup python -u train.py --dataroot yourdata --name tot1ce_pre_trained --gpu_ids 7 --batchSize 100 --model resvit_many --which_model_netG res_cnn  --which_direction AtoB --lambda_A 100 --dataset_mode aligned --norm batch --pool_size 0 --output_nc 1 --input_nc 3 --loadSize 128 --fineSize 128 --niter 50 --niter_decay 50 --save_epoch_freq 5 --checkpoints_dir checkpoints/ --display_id 0 --lr 0.0002 >log_tot1ce_pre.log 2>&1&

根据自己的实际情况调节dataroot、gpu、batchSize