src/data_loader_han.py

"""
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'])