
import os
import numpy as np
import tensorflow as tf
from image_quality.utils import utils


class TrainDataGenerator(tf.keras.utils.Sequence):
  '''inherits from Keras Sequence base object, allows to use multiprocessing in .fit_generator'''
  def __init__(self, samples, img_dir, batch_size, n_classes, basenet_preprocess,
         img_load_dims=(256, 256), img_crop_dims=(224, 224), shuffle=True):
    self.samples = samples
    self.img_dir = img_dir
    self.batch_size = batch_size
    self.n_classes = n_classes
    self.basenet_preprocess = basenet_preprocess  # Keras basenet specific preprocessing function
    self.img_load_dims = img_load_dims  # dimensions that images get resized into when loaded
    self.img_crop_dims = img_crop_dims  # dimensions that images get randomly cropped to
    self.shuffle = shuffle
    self.on_epoch_end()  # call ensures that samples are shuffled in first epoch if shuffle is set to True

  def __len__(self):
    return int(np.ceil(len(self.samples) / self.batch_size))  # number of batches per epoch

  def __getitem__(self, index):
    batch_indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]  # get batch indexes
    batch_samples = [self.samples[i] for i in batch_indexes]  # get batch samples
    X, y = self.__data_generator(batch_samples)
    return X, y

  def on_epoch_end(self):
    self.indexes = np.arange(len(self.samples))
    if self.shuffle is True:
      np.random.shuffle(self.indexes)

  def __data_generator(self, batch_samples):
    # initialize images and labels tensors for faster processing
    X = np.empty((len(batch_samples), *self.img_crop_dims, 3))
    y = np.empty((len(batch_samples), self.n_classes))

    for i, sample in enumerate(batch_samples):
      # load and randomly augment image
      img_file = os.path.join(self.img_dir, '{}'.format(sample['image_id']))
      img = utils.load_image(img_file, self.img_load_dims)
      if img is not None:
        img = utils.random_crop(img, self.img_crop_dims)
        img = utils.random_horizontal_flip(img)
        X[i, ] = img

      # normalize labels
      y[i, ] = utils.normalize_labels(sample['label'])

    # apply basenet specific preprocessing
    # input is 4D numpy array of RGB values within [0, 255]
    X = self.basenet_preprocess(X)

    return X, y


class TestDataGenerator(tf.keras.utils.Sequence):
  '''inherits from Keras Sequence base object, allows to use multiprocessing in .fit_generator'''
  def __init__(self, samples, img_dir, batch_size, n_classes, basenet_preprocess,
         img_load_dims=(224, 224)):
    self.samples = samples
    self.img_dir = img_dir
    self.batch_size = batch_size
    self.n_classes = n_classes
    self.basenet_preprocess = basenet_preprocess  # Keras basenet specific preprocessing function
    self.img_load_dims = img_load_dims  # dimensions that images get resized into when loaded
    self.on_epoch_end()  # call ensures that samples are shuffled in first epoch if shuffle is set to True

  def __len__(self):
    return int(np.ceil(len(self.samples) / self.batch_size))  # number of batches per epoch

  def __getitem__(self, index):
    batch_indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]  # get batch indexes
    batch_samples = [self.samples[i] for i in batch_indexes]  # get batch samples
    X, y = self.__data_generator(batch_samples)
    return X, y

  def on_epoch_end(self):
    self.indexes = np.arange(len(self.samples))

  def __data_generator(self, batch_samples):
    # initialize images and labels tensors for faster processing
    X = np.empty((len(batch_samples), *self.img_load_dims, 3))
    y = np.empty((len(batch_samples), self.n_classes))

    for i, sample in enumerate(batch_samples):
      # load and randomly augment image
      img_file = os.path.join(self.img_dir, '{}'.format(sample['image_id']))
      img = utils.load_image(img_file, self.img_load_dims)
      if img is not None:
        X[i, ] = img

      # normalize labels
      if sample.get('label') is not None:
        y[i, ] = utils.normalize_labels(sample['label'])

    # apply basenet specific preprocessing
    # input is 4D numpy array of RGB values within [0, 255]
    X = self.basenet_preprocess(X)

    return X, y

[((1440, 1471), 'numpy.random.shuffle', 'np.random.shuffle', (['self.indexes'], {}), '(self.indexes)\n', (1457, 1471), True, 'import numpy as np\n'), ((1878, 1924), 'image_quality.utils.utils.load_image', 'utils.load_image', (['img_file', 'self.img_load_dims'], {}), '(img_file, self.img_load_dims)\n', (1894, 1924), False, 'from image_quality.utils import utils\n'), ((2118, 2157), 'image_quality.utils.utils.normalize_labels', 'utils.normalize_labels', (["sample['label']"], {}), "(sample['label'])\n", (2140, 2157), False, 'from image_quality.utils import utils\n'), ((3876, 3922), 'image_quality.utils.utils.load_image', 'utils.load_image', (['img_file', 'self.img_load_dims'], {}), '(img_file, self.img_load_dims)\n', (3892, 3922), False, 'from image_quality.utils import utils\n'), ((1965, 2007), 'image_quality.utils.utils.random_crop', 'utils.random_crop', (['img', 'self.img_crop_dims'], {}), '(img, self.img_crop_dims)\n', (1982, 2007), False, 'from image_quality.utils import utils\n'), ((2022, 2055), 'image_quality.utils.utils.random_horizontal_flip', 'utils.random_horizontal_flip', (['img'], {}), '(img)\n', (2050, 2055), False, 'from image_quality.utils import utils\n'), ((4055, 4094), 'image_quality.utils.utils.normalize_labels', 'utils.normalize_labels', (["sample['label']"], {}), "(sample['label'])\n", (4077, 4094), False, 'from image_quality.utils import utils\n')]