self-attention-cv：用于三维分割的子体积采样训练过程的演示

原文在此： self-attention-cv.

数据集在此： dataset.

from google.colab import drive
drive.mount('/gdrive')
import zipfile
root_path = '/gdrive/My Drive/DATASETS/iSeg-2019-Training.zip' 
!echo "Download and extracting folders..."
zip_ref = zipfile.ZipFile(root_path, 'r')
zip_ref.extractall("./")
zip_ref.close()
!echo "Finished"
!pip install torchio
!pip install self-attention-cv

训练示例

import glob
import torchio as tio
import torch
from torch.utils.data import DataLoader

paths_t1 = sorted(glob.glob('./iSeg-2019-Training/*T1.img'))
paths_t2 = sorted(glob.glob('./iSeg-2019-Training/*T2.img'))
paths_seg = sorted(glob.glob('./iSeg-2019-Training/*label.img'))
assert len(paths_t1) == len(paths_t2) == len(paths_seg)

subject_list = []
for pat in zip(paths_t1, paths_t2, paths_seg):
  path_t1, path_t2, path_seg = pat
  subject = tio.Subject(t1=tio.ScalarImage(path_t1,),
              t2=tio.ScalarImage(path_t2,),
              label=tio.LabelMap(path_seg)) 
  subject_list.append(subject)


transforms = [tio.RescaleIntensity((0, 1)),tio.RandomAffine() ]
transform = tio.Compose(transforms)

subjects_dataset = tio.SubjectsDataset(subject_list, transform=transform)

patch_size = 24
queue_length = 300
samples_per_volume = 50
sampler = tio.data.UniformSampler(patch_size)

patches_queue = tio.Queue(
subjects_dataset,
queue_length,
samples_per_volume,sampler, num_workers=1)

patches_loader = DataLoader(patches_queue, batch_size=16)

from self_attention_cv.Transformer3Dsegmentation import Transformer3dSeg

def crop_target(img, target_size):
  dim = img.shape[-1]
  center = dim//2
  start_dim = center - (target_size//2) - 1
  end_dim = center + (target_size//2)
  return img[:,0,start_dim:end_dim,start_dim:end_dim,start_dim:end_dim].long()

target_size = 3 # as in the paper 
patch_dim = 8
num_epochs = 50
num_classes = 4
model = Transformer3dSeg(subvol_dim=patch_size, patch_dim=patch_dim,
                         in_channels=2, blocks=2, num_classes=num_classes).cuda()

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
print(len(patches_loader))
for epoch_index in range(num_epochs):
  epoch_loss = 0
  for c,patches_batch in enumerate(patches_loader):
    optimizer.zero_grad()
    
    input_t1 = patches_batch['t1'][tio.DATA]  
    input_t2 = patches_batch['t2'][tio.DATA]

    input_tensor = torch.cat([input_t1, input_t2], dim=1).cuda()
    
    
    logits = model(input_tensor) # 8x8x8 the 3d transformer-based approach

    # for the 3d transformer-based approach the target must be cropped again to the desired size
    targets = patches_batch['label'][tio.DATA]  
    
    cropped_target = crop_target(targets, target_size).cuda()

    loss = criterion(logits, cropped_target)
    loss.backward()
    optimizer.step()
    epoch_loss = epoch_loss+loss.cpu().item()

  print(f'epoch {epoch_index} loss {epoch_loss/c}')

结果

import torch
import torch.nn as nn
import torchio as tio
patch_overlap = 0
patch_size = 24, 24, 24
target_patch_size = 3

#input sampling
grid_sampler = tio.inference.GridSampler(subject_list[0], patch_size, patch_overlap)
patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=4)
# target vol sampling
grid_sampler_target = tio.inference.GridSampler(subject_list[0], target_patch_size, patch_overlap)
aggregator = tio.inference.GridAggregator(grid_sampler_target)
target_loader = torch.utils.data.DataLoader(grid_sampler_target, batch_size=4)

model.eval()

with torch.no_grad():
  for patches_batch,target_patches in zip(patch_loader,target_loader):

    input_t1 = patches_batch['t1'][tio.DATA]  
    input_t2 = patches_batch['t2'][tio.DATA]
    input_tensor = torch.cat([input_t1, input_t2], dim=1).float().cuda()

    locations = target_patches[tio.LOCATION]
    logits = model(input_tensor)
    labels = logits.argmax(dim=tio.CHANNELS_DIMENSION, keepdim=True)
    outputs = labels
    aggregator.add_batch(outputs.type(torch.int32), locations)

  print('output tensor shape:',outputs.shape)
  output_tensor = aggregator.get_output_tensor()
  print(output_tensor.shape)