Iterate on DataGenerator loops indefinitely

Hello everybody,

I implemented a DataGenerator

class DataGenerator(Sequence):

def __init__(self, df, augmentation = None, policy = None , batch_size=32, dim=(128, 128), n_channels=1, shuffle=False):
    self.data_df =  255 - df.iloc[:, 1:].values.reshape(-1, HEIGHT, WIDTH).astype(np.float64)
    self.imgs_id_df = df.iloc[:,0]
    self.batch_size=batch_size
    self.dim = dim
    self.n_channels = n_channels
    self.shuffle = shuffle
    self.augment = augmentation
    self.policy = policy
    self.n_channels = n_channels
    self.on_epoch_end()

def __len__(self):
     return int(np.floor(len(self.data_df) / self.batch_size))

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

    '''Find list of Images''' 
    data_df_temp = [self.data_df[k] for k in indexes]
    '''Generate list of Images ID'''
    imgs_id_df_temp = [self.imgs_id_df[k] for k in indexes]

    imgs_id = np.array(imgs_id_df_temp)
    '''Generate data'''
    X = self._generate_X(data_df_temp)

    return X, imgs_id

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

def _generate_X(self, data_df_temp):
    '''Initialization'''
    X = np.empty((self.batch_size, *self.dim, self.n_channels))
    '''Generate data'''
    for idx in range(len(data_df_temp)):
        '''Store sample'''
        img = (data_df_temp[idx]*(255.0/data_df_temp[idx].max())).astype(np.float64)
        img = self._load_grayscale_image(img)

        img = pd.DataFrame(img)
        img = img.values.reshape(-1, SIZE, SIZE, N_CHANNELS)
        X[idx,] = img
    return X

def _load_grayscale_image(self, img):
    img = crop_resize(img)
    return img

when i am iterating using for loop below it is going into loops indefinitely for x, y in train_generator:

test_generator = DataGenerator(test_df, policy = None, batch_size=32, augmentation = ImageDataGenerator( horizontal_flip=False, vertical_flip=False, )) for X_test, imgId in test_generator: pred = model.predict(X_test)

keras-team / keras-preprocessing

Iterate on DataGenerator loops indefinitely #269

I implemented a DataGenerator