原文在此: 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)