File size: 4,559 Bytes
9759eef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
Data-loading helpers for the HaN dose-prediction models.
Trimmed from the original research toolkit: only the functions the released
data loader and inference script actually call are kept, which drops the
cv2 / matplotlib / SimpleITK plotting dependencies.
Derived from the data pipeline of:
Riqiang Gao, Bin Lou, Zhoubing Xu, Dorin Comaniciu, Ali Kamen.
"Flexible-CM GAN: Towards Precise 3D Dose Prediction in Radiotherapy." CVPR 2023.
"""
import os
import numpy as np
import torch
import nibabel as nib
from monai.transforms import (
Compose,
Resized,
RandFlipd,
RandRotated,
SpatialPadd,
SpatialCropd,
RandSpatialCropd,
)
def combine_oar(tmp_dict, need_list,norm_oar = True, OAR_DICT = None):
'''
this function is used to support the data loader.
tmp_dict: the dictionary of the data loaded from the npz file
need_list: the list of the OARs needed to be combined
norm_oar: if True, the OARs will be normalized to the same scale
OAR_DICT: the dictionary of the OARs, the key is the name of the OAR, the value is the index of the OAR in the combined
'''
comb_oar = torch.zeros(tmp_dict['img'].shape)
cat_oar = torch.zeros([len(OAR_DICT)] + list(tmp_dict['img'].shape)[1:])
for key in OAR_DICT.keys():
if key not in need_list:
continue
if key in tmp_dict.keys():
single_oar = tmp_dict[key]
else:
single_oar = torch.zeros(tmp_dict['img'].shape)
cat_oar[OAR_DICT[key]-1: OAR_DICT[key]] = single_oar
if norm_oar:
comb_oar = torch.maximum(comb_oar, single_oar.round() * (1.0 + 4.0 * OAR_DICT[key] / 30))
else:
comb_oar = torch.maximum(comb_oar, single_oar.round() * OAR_DICT[key]) # changed in this version
return comb_oar, cat_oar
def combine_ptv(tmp_dict, scaled_dose_dict):
prescribed_dose = []
cat_ptv = torch.zeros([3] + list(tmp_dict['img'].shape)[1:])
prescribed_dose = [0] * 3
comb_ptv = torch.zeros(tmp_dict['img'].shape)
cnt = 0
for key in scaled_dose_dict.keys():
tmp_ptv = tmp_dict[key] * scaled_dose_dict[key]
prescribed_dose[cnt] = scaled_dose_dict[key]
cat_ptv[cnt] = tmp_ptv
comb_ptv = torch.maximum(comb_ptv, tmp_ptv)
cnt += 1
# sort the cat_ptv according to the prescribed dose
paired = [(cat_ptv[i], prescribed_dose[i]) for i in range(len(prescribed_dose))]
paired_sorted = sorted(paired, key=lambda x: x[1], reverse=True)
cat_ptv = torch.stack([x[0] for x in paired_sorted])
prescribed_dose = [x[1] for x in paired_sorted]
return comb_ptv, prescribed_dose, cat_ptv
def tr_augmentation(KEYS, in_size, out_size, crop_center):
return Compose([
SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], allow_missing_keys = True),
SpatialPadd(keys = KEYS, spatial_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], mode = 'constant', allow_missing_keys = True),
RandSpatialCropd(keys = KEYS, roi_size = [int(in_size[0] * 0.85), int(in_size[1] * 0.85), int(in_size[2] * 0.85)], max_roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], random_center = True, random_size = True, allow_missing_keys = True),
RandRotated(keys = KEYS, prob=0.8, range_x= 1, range_y = 0.2, range_z = 0.2, allow_missing_keys = True),
RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 0, allow_missing_keys = True),
RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 1, allow_missing_keys = True),
RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 2, allow_missing_keys = True),
Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True),
])
def tt_augmentation(KEYS, in_size, out_size, crop_center):
return Compose([
SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = in_size, allow_missing_keys = True),
SpatialPadd(keys = KEYS, spatial_size = in_size, mode = 'constant', allow_missing_keys = True),
Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True),
])
def load_nifti(path):
"""
Load a .nii or .nii.gz file and return a numpy array.
torch shape: [B, C, Z, Y, X] where B=1, C=1, Z=depth, Y=height, X=width
"""
nii = nib.load(path)
return np.transpose(nii.get_fdata(), (2, 1, 0)).astype(np.float32)
|