load data layer demo

[JunMa11]

I want to use nifit data to train a unet, here is a demo for datalayer. Anyone who finds mistakes please make a suggestion. Thanks

import numpy as np
import nibabel as nb

class DataLayer():
    def __init__(self, opt):
        # 'batch_size': 2,
        # 'learning_rate': 0.0001,
        # 'lr_decay': 0.5,
        # 'save_model_every': 100,
        # 'checkpoint_path': 'checkpoints/unet',
        # 'epoch': 5,
        # 'load_from_checkpoint': 'unet-654'
        self.batch_size = opt['batch_size']
        self.learning_rate = opt['learning_rate']
        self.lr_decay = opt['lr_decay']
        self.save_model_every = opt['save_model_every']
        self.checkpoint_path = opt['checkpoint_path']
        self.epoch = opt['epoch']
        self.load_from_checkpoint = opt['load_from_checkpoint']

    def get_iter_epoch(self):
        return self.epoch

    def load_batch(self):
        # read custome data
        TrainFilePath = "H:/e_PublicData/H_LiverTumor2017/Training Batch 2/" #data pyth
        Train_Case = 120
        data_name = TrainFilePath + 'volume-' + str(Train_Case) + '.nii'
        label_name = TrainFilePath + 'segmentation-' + str(Train_Case) + '.nii'
        Train_nii = nb.load(data_name) # load data
        Train_Label_nii = nb.load(label_name) #load label
        Train_data = Train_nii.get_data()
        Train_data = Train_data.astype(np.uint8)
        Train_Label_data = Train_Label_nii.get_data()
        Train_Label_data[Train_Label_data == 2] = 1 # laber number:2  background 0, ground truth 1
        [nx, ny, nz] = Train_data.shape
        X = np.zeros([nz, nx, ny, 1], dtype= float)
        Y = np.zeros([nz, nx, ny, 1], dtype= float)
        Train_data = Train_data.transpose(2, 0, 1)
        Train_Label_data = Train_Label_data.transpose(2, 0, 1)
        X[..., 0] = Train_data
        Y[..., 0] = Train_Label_data
        X_temp = X[0:2, ...]
        Y_temp = Y[0:2, ...]
        return X_temp, Y_temp

[zizhaozhang]

Hi I have updated the code to give an example of loader