Brain Tumor Detection-CV Data Generator data acquisition

Don't talk nonsense and go directly to the code (modified classification code)

import tensorflow
class DataGenerator(tensorflow.keras.utils.Sequence):
    'Generates data for Keras'
    def __init__(self, list_IDs, batch_size=32, dim=(224,224), n_channels=3,
                 n_classes=2, shuffle=True):
        'Initialization'
        self.dim = dim
        self.batch_size = batch_size
        self.list_IDs = list_IDs
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        # Generate indexes of the batch
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]

        # Find list of IDs
        list_IDs_temp = [self.list_IDs[k] for k in indexes]

        # Generate data
        X, y = self.__data_generation(list_IDs_temp)
        return X, y

    def on_epoch_end(self):
        'Updates indexes after each epoch'
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __data_generation(self, list_IDs_temp):
        'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
        # Initialization
        X = np.empty(shape=(self.batch_size, *self.dim, self.n_channels))
        y = np.empty(shape=(self.batch_size, *self.dim))
        # Generate data
        for i, ID in enumerate(list_IDs_temp):
            # Store sample
            img1 = np.load(ID)[:,:,:]
            seg1 = np.load(mask[i])[:,:]
            # Add data augmentation here
            image = augmentation(img=img1, seg=seg1)
            X[i] = image["image"]
            y[i] = image["mask"]
            X[i] = add_noise(X[i])

            # Store class
        return X, y

The first three paragraphs are similar to the Dataset class, providing constructor, length and sample generation. On epoch end means that after one epoch ends, the existing indexes are shuffled, and several new batch_size data are randomly obtained as an epoch. data_generation runs first, reads the path in img as numpy, and randomly crops it to a size of 224*224, adds noise or flips it with probability, stores it as X and y, and throws it to the model for training. (data augmentation). X is the 3-channel RGB image img, and y is the single-channel mask.

Generator will return a tuple (tuple) containing two numpy arrays img: (32, 224, 224, 3) and mask: (32, 224, 224).

Compared with classification, the value of y in segmentation should be the mask corresponding to the ID (that is, img), not the label 0 or 1. At the same time, it should be noted that img has three channels, and a 3 (channel) should be added to the dimension, and seg has only one channel, so there is no need to add it. Since all pictures are randomly cropped 240 → 224, the array generated in the generator should also adapt to the size.