""" Data loader for the head-and-neck (HaN) dose-prediction models. Produces the 5-channel input the released HaN checkpoints expect: data_dict['data'] : [B, 5, 128, 192, 192] (Z, Y, X) ch0 comb_optptv - PTVs weighted by prescription dose (optimisation set) ch1 comb_ptv - PTVs weighted by prescription dose ch2 comb_oar - all OARs merged into one label-encoded channel ch3 body - external/BODY mask ch4 img - CT, clipped to [down_HU, up_HU] and divided by denom_norm_HU data_dict['label'] : [B, 1, 128, 192, 192] dose, D97-normalised and divided by `dose_div_factor`, masked by BODY The *pretrained* (source-domain) checkpoints instead take 8 channels: the five above plus beam_plate, angle_plate and prompt_extend. Those three are all-zero for this dataset, which is why the released transfer-learned models drop them. See README.md for how the input stem was re-initialised. Data pipeline derived from: Riqiang Gao, Bin Lou, Zhoubing Xu, Dorin Comaniciu, and Ali Kamen. "Flexible-CM GAN: Towards Precise 3D Dose Prediction in Radiotherapy." CVPR 2023. """ from torch.utils.data import Dataset, DataLoader import pandas as pd import torch import numpy as np import json import yaml from scipy import ndimage from toolkit_han import * HaN_OAR_LIST = ["Bone_Mandible", "Brainstem", "Cavity_Oral", "Esophagus", "Glnd_Thyroid", "Larynx", "Parotid_L", "Parotid_R", "SpinalCord", "Submandibular_L", "Submandibular_R", "Brain", "Eye_L", "Eye_R", "OpticNrv_L", "OpticNrv_R", "OpticChiasm", "Lens_L", "Lens_R", "Lips", "Trachea", "Lung_L", "SpinalCord_PRV05", ] # 23 OARs in HaN HaN_OAR_DICT = {HaN_OAR_LIST[i]: (i+1) for i in range(len(HaN_OAR_LIST))} class MyDataset(Dataset): def __init__(self, cfig, phase): ''' phase: train, validation, or testing cfig: the configuration dictionary train_bs: training batch size val_bs: validation batch size num_workers: the number of workers when call the DataLoader of PyTorch csv_root: the meta data file, include patient id and some conditions of the plan. scale_dose_dict: path of a dictionary. The dictionary includes the prescribed doses of the PTVs. pat_obj_dict: path of a dictionary. The dictionary includes the OARs names used in optimization. down_HU: bottom clip of the CT HU value. up_HU: upper clip of the CT HU value. denom_norm_HU: the denominator when normalizing the CT. in_size & out_size: the size parameters used in data transformation. norm_oar: True or False. Normalize the OAR channel or not. CatStructures: True or False. Concat the PTVs and OARs in multiple channels, or merge them in one channel, respectively. dose_div_factor: the value used to normalize dose. ''' self.cfig = cfig df = pd.read_csv(cfig['csv_root']) df = df.loc[df['dev_split'] == phase] # !!! Actual train + valid + test split. self.phase = phase self.data_list = df['PatientID'].tolist() self.site_list = df['site'].tolist() self.cohort_list = df['cohort'].tolist() self.scale_dose_Dict = json.load(open(cfig['scale_dose_dict'], 'r')) self.pat_obj_dict = json.load(open(cfig['pat_obj_dict'], 'r')) def __len__(self): return len(self.data_list) def __getitem__(self, index): ID = self.data_list[index] PatientID = ID if self.phase == 'test': root_dir = os.path.join(self.cfig['data_root'], 'test') else: root_dir = os.path.join(self.cfig['data_root'], 'train_valid') patient_dir = os.path.join(root_dir, ID) ct = load_nifti(os.path.join(patient_dir, 'CT.nii.gz')) dose = load_nifti(os.path.join(patient_dir, 'dose.nii.gz')) body = load_nifti(os.path.join(patient_dir, 'Masks', 'BODY.nii.gz')) # print(f"Loading patient {ID} with CT shape {ct.shape} and dose shape {dose.shape}") # Normalize CT img. ct = np.clip(ct, self.cfig['down_HU'], self.cfig['up_HU']) / self.cfig['denom_norm_HU'] ori_img_size = ct.shape # Normalize dose by D97 if 'PTV_High' in self.scale_dose_Dict[PatientID]: ptv_high = self.scale_dose_Dict[PatientID]['PTV_High'] pdose = ptv_high['PDose'] struct_name = ptv_high['StructName'] ptv_mask = load_nifti(os.path.join(patient_dir, 'Masks', 'PTV', f"{struct_name}.nii.gz")) d97 = np.percentile(dose[ptv_mask > 0], 3) norm_scale = pdose / (d97 + 1e-5) dose = dose * norm_scale / self.cfig['dose_div_factor'] dose = np.clip(dose, 0, pdose * 1.2) isocenter = np.array(ct.shape) // 2 # Load masks if self.site_list[index] < 1.5: OAR_LIST = HaN_OAR_LIST OAR_DICT = HaN_OAR_DICT else: raise NotImplementedError("Only HaN site supported in this loader.") try: need_list = self.pat_obj_dict[ID.split('+')[0]] except: need_list = OAR_LIST In_dict = { 'img': ct, 'body': body, 'dose': dose, 'isocenter': isocenter, } for name in need_list: path = os.path.join(patient_dir, 'Masks', 'OAR', f"{name}.nii.gz") if os.path.exists(path): In_dict[name] = load_nifti(path) ptv_dict = self.scale_dose_Dict[PatientID] opt_dose_dict = {} dose_dict = {} for k in ['PTV_High', 'PTV_Mid', 'PTV_Low']: if k in ptv_dict: pdose = ptv_dict[k]['PDose'] / self.cfig['dose_div_factor'] struct = ptv_dict[k]['StructName'] path = os.path.join(patient_dir, 'Masks', 'PTV', f"{struct}.nii.gz") if os.path.exists(path): In_dict[struct] = load_nifti(path) opt_dose_dict[struct] = pdose dose_dict[struct] = pdose # Convert to torch tensor KEYS = list(In_dict.keys()) for key in list(KEYS): if isinstance(In_dict[key], np.ndarray) and len(In_dict[key].shape) == 3: In_dict[key] = torch.from_numpy(In_dict[key].astype('float'))[None] else: KEYS.remove(key) use_aug = self.cfig.get('with_aug', True) if self.phase == 'train': if use_aug: self.aug = tr_augmentation(KEYS, self.cfig['in_size'], self.cfig['out_size'], isocenter) else: self.aug = tt_augmentation(KEYS, self.cfig['in_size'], self.cfig['out_size'], isocenter) else: self.aug = tt_augmentation(KEYS, self.cfig['in_size'], self.cfig['out_size'], isocenter) In_dict = self.aug(In_dict) comb_oar, cat_oar = combine_oar(In_dict, need_list, self.cfig['norm_oar'], OAR_DICT) comb_optptv, prs_opt, cat_optptv = combine_ptv(In_dict, opt_dose_dict) comb_ptv, _, cat_ptv = combine_ptv(In_dict, dose_dict) # Plan-level metadata. The source-domain (8-channel) models consumed this # broadcast into a `prompt_extend` channel; the released 5-channel models # do not. It is still returned in data_dict for reference. See channels.py. prompt = torch.tensor([1.0, len(prs_opt), self.site_list[index], self.cohort_list[index]]).float() # 5-channel input: comb_optptv, comb_ptv, comb_oar, body, img. # The source-domain models additionally took beam_plate, angle_plate and # prompt_extend; this dataset has no beam geometry, so those three # channels were dropped rather than fed as zeros, and the input stem was # re-initialised. See README.md and channels.py. if self.cfig['CatStructures']: input_data = torch.cat((cat_optptv, cat_ptv, cat_oar, In_dict['body'], In_dict['img']), dim=0) else: input_data = torch.cat((comb_optptv, comb_ptv, comb_oar, In_dict['body'], In_dict['img']), dim=0) data_dict = { 'data': input_data, 'label': In_dict['dose'] * In_dict['body'], 'oar': cat_oar, 'ptv': cat_ptv, 'optptv': cat_optptv, 'id': ID, 'ori_img_size': torch.tensor(ori_img_size), 'ori_isocenter': torch.tensor(isocenter), 'prompt': prompt } prescribed_dose = [ptv_dict['PTV_High']['PDose'] if 'PTV_High' in ptv_dict else 0] prescribed_dose.append(ptv_dict['PTV_Mid']['PDose'] if 'PTV_Mid' in ptv_dict else 0) prescribed_dose.append(ptv_dict['PTV_Low']['PDose'] if 'PTV_Low' in ptv_dict else 0) data_dict['prescrbed_dose'] = torch.tensor(prescribed_dose).float() return data_dict class GetLoader(object): def __init__(self, cfig): super().__init__() self.cfig = cfig def train_dataloader(self): dataset = MyDataset(self.cfig, phase='train') return DataLoader(dataset, batch_size=self.cfig['train_bs'], shuffle=True, num_workers=self.cfig['num_workers']) def val_dataloader(self): dataset = MyDataset(self.cfig, phase='valid') return DataLoader(dataset, batch_size=self.cfig['val_bs'], shuffle=False, num_workers=self.cfig['num_workers']) def test_dataloader(self): dataset = MyDataset(self.cfig, phase='test') return DataLoader(dataset, batch_size=self.cfig['val_bs'], shuffle=False, num_workers=self.cfig['num_workers'])