src/data_loader_pancreas.py

"""
Data loader for the pancreas (PancVMAT) dose-prediction model.

Same 5-channel layout as data_loader_han.py, with two deliberate differences:

  * `combine_optptv` (see toolkit_pancreas.py) subtracts every overlapping OAR
    from the PTV before building channel 0.
  * The dose target is NOT D97-renormalised (`norm_scale = 1.0`); the planned
    dose is used as-is, only divided by `dose_div_factor`.

The OAR list is the 14 PancVMAT structures below, not the 23 HaN ones.

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_pancreas import *

Panc_OAR_LIST = ["Kidney_L.nii.gz",
                "Kidney_R.nii.gz",
                "Kidneys.nii.gz",
                "Liver.nii.gz",
                "SpinalCord",
                "SpinalCanal.nii.gz",
                "Esophagus.nii.gz",
                "Heart.nii.gz",

                "Stomach.nii.gz",
                "Stomach_PRV.nii.gz",
                "Bowel_Small.nii.gz",
                "Bowel_Small_PRV.nii.gz",
                "Bowel_Large.nii.gz",
                "Bowel_Large_PRV.nii.gz",
                ]  # 14 OARs in PancVMAT

Panc_OAR_DICT = {Panc_OAR_LIST[i]: (i+1) for i in range(len(Panc_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'))
        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
        
        # NO Normalize dose by D95!!!
        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"))

            # d95 = np.percentile(dose[ptv_mask > 0], 5)
            # norm_scale = pdose / (d95 + 1e-5)
            norm_scale = 1.0
            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] == 2:
            OAR_LIST = Panc_OAR_LIST
            OAR_DICT = Panc_OAR_DICT
        else:
            raise NotImplementedError("Only PancVMAT 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_optptv(In_dict, opt_dose_dict, OAR_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.
        # Channels 0-2 are computed differently from the HaN loader -- see
        # channels.py (PANCREAS_5CH). Do not mix the two encodings.
        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'])