__getitem__ for data generator.py

def __getitem__(self, index):
      data_index_min = int(index*self.batch_size)
      data_index_max = int(min((index+1)*self.batch_size, len(self.image_filenames)))

      indexes = self.image_filenames[data_index_min:data_index_max]
      this_batch_size = len(indexes) # The last batch can be smaller than the others

      X = np.empty((this_batch_size, self.image_size, self.image_size, 3), dtype=np.float32)
      y = np.empty((this_batch_size, self.image_size, self.image_size, self.nb_y_features), dtype=np.uint8)

      for i, sample_index in enumerate(indexes):

          X_sample, y_sample = self.read_image_mask(self.image_filenames[index * self.batch_size + i], 
                                                  self.mask_names[index * self.batch_size + i])

          # if augmentation is defined, we assume its a train set
          if self.augmentation is not None:

              # Augmentation code
              augmented = self.augmentation(self.image_size)(image=X_sample, mask=y_sample)
              image_augm = augmented['image']
              mask_augm = augmented['mask'].reshape(self.image_size, self.image_size, self.nb_y_features)
              # divide by 255 to normalize images from 0 to 1
              X[i, ...] = image_augm/255
              y[i, ...] = mask_augm
          else:
              ...
  return X, y