def pre_processing(dict_images):
MR = dict_images['MR']
MR = np.clip(MR / 2048, a_min=0, a_max=1)
Mask = dict_images['Mask']
_, D, H, W = MR.shape
heatmap_generator = HeatmapGenerator(image_size=(D, H, W),
sigma=2.,
scale_factor=1.,
normalize=True,
size_sigma_factor=8,
sigma_scale_factor=2,
dtype=np.float32)
list_landmarks = dict_images['list_landmarks']
index = random.randint(10, 18)
while True in np.isnan(list_landmarks[index]):
index = random.randint(10, 18)
heatmap = heatmap_generator.generate_heatmap(landmark=list_landmarks[index])[np.newaxis, :, :, :]
Mask = np.where(Mask == index + 1, 1, 0) # just segment one IVD
if D > 12:
start = random.choice([i for i in range(D - 12 + 1)])
MR = crop(MR, start=start, end=start + 12, axis='z')
heatmap = crop(heatmap, start=start, end=start + 12, axis='z')
Mask = crop(Mask, start=start, end=start + 12, axis='z')
MR = crop_to_center(MR, list_landmarks[index], dsize=(12, 64, 96))
heatmap = crop_to_center(heatmap, list_landmarks[index], dsize=(12, 64, 96))
Mask = crop_to_center(Mask, list_landmarks[index], dsize=(12, 64, 96))
return [np.concatenate((MR, heatmap)), Mask]
def inference(trainer, list_case_dirs, save_path, do_TTA=False):
if not os.path.exists(save_path):
os.mkdir(save_path)
if do_TTA:
TTA_mode = [[], [2], [4], [2, 4]]
else:
TTA_mode = [[]]
with torch.no_grad():
trainer.setting.network.eval()
for case_dir in tqdm(list_case_dirs):
assert os.path.exists(case_dir), case_dir + 'does not exist!'
case_id = case_dir.split('/')[-1]
dict_images = read_data(case_dir)
list_images = pre_processing(dict_images)
MR = list_images[0]
# MR = torch.from_numpy(MR)
list_IVD_landmarks = list_images[1]
C, D, H, W = MR.shape
dsize = (12, 64, 96)
# all pred_IVDMask will be insert into this tensor
pred_Mask = torch.zeros(C, D, H, W).to(trainer.setting.device)
heatmap_generator = HeatmapGenerator(image_size=(D, H, W),
sigma=2.,
scale_factor=1.,
normalize=True,
size_sigma_factor=8,
sigma_scale_factor=2,
dtype=np.float32)
for index, landmark in enumerate(list_IVD_landmarks):
if True in np.isnan(landmark):
continue
temp = torch.zeros(C, D, H, W).to(trainer.setting.device)
heatmap = heatmap_generator.generate_heatmap(landmark)[
np.newaxis, :, :, :] # (1, D, H, W)
# heatmap = torch.from_numpy(heatmap)
input_ = np.concatenate((MR, heatmap), axis=0) # (2, D, H, W)
if D > 12:
input_, patch, pad = crop_to_center(input_,
landmark=landmark,
dsize=dsize)
input_ = np.stack(
(input_[:, :12, :, :], input_[:, -12:, :, :]),
axis=0) # (2, 2, 12, H, W)
input_ = torch.from_numpy(input_).to(
trainer.setting.device)
# pred_IVDMask = trainer.setting.network(input_) # (2, 2, 12, 128, 128)
pred_IVDMask = test_time_augmentation(
trainer, input_, TTA_mode)
pred_IVDMask = post_processing(
pred_IVDMask, D,
device=trainer.setting.device) # (1, 2, D, 128, 128)
pred_IVDMask = nn.Softmax(dim=1)(pred_IVDMask)
pred_IVDMask = torch.argmax(pred_IVDMask,
dim=1) # (1, D, 128, 128)
else:
input_, patch, pad = crop_to_center(input_,
landmark=landmark,
dsize=dsize)
input_ = torch.from_numpy(input_).unsqueeze(0).to(
trainer.setting.device)
# pred_IVDMask = trainer.setting.network(input_) # (1, 2, 12, 128, 128)
pred_IVDMask = test_time_augmentation(
trainer, input_, TTA_mode)
pred_IVDMask = nn.Softmax(dim=1)(pred_IVDMask)
pred_IVDMask = torch.argmax(pred_IVDMask,
dim=1) # (1, 12, 128, 128)
bh, eh, bw, ew = patch
pad_h_1, pad_h_2, pad_w_1, pad_w_2 = pad
if pad_h_1 > 0:
pred_IVDMask = pred_IVDMask[:, :, pad_h_1:, :]
if pad_h_2 > 0:
pred_IVDMask = pred_IVDMask[:, :, :-pad_h_2, :]
if pad_w_1 > 0:
pred_IVDMask = pred_IVDMask[:, :, :, pad_w_1:]
if pad_w_2 > 0:
pred_IVDMask = pred_IVDMask[:, :, :, :-pad_w_2]
pred_IVDMask = torch.where(pred_IVDMask > 0, index + 11, 0)
temp[:, :, bh:eh, bw:ew] = pred_IVDMask
pred_Mask += temp
pred_Mask = pred_Mask.cpu().numpy() # (1, 12, 128, 128)
# Save prediction to nii image
template_nii = sitk.ReadImage(case_dir + '/MR_512.nii.gz')
prediction_nii = sitk.GetImageFromArray(pred_Mask[0])
prediction_nii = copy_sitk_imageinfo(template_nii, prediction_nii)
if not os.path.exists(save_path + '/' + case_id):
os.mkdir(save_path + '/' + case_id)
sitk.WriteImage(prediction_nii,
save_path + '/' + case_id + '/pred_IVDMask.nii.gz')