Trial of a DALIDataset with external input for Keras

keras_pipeline_trial_fixed.py

import glob
import numpy as np
import cupy as cp
import imageio
from random import shuffle
from nvidia.dali import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.plugin.tf as dali_tf
import tensorflow as tf

BATCH_SIZE = 16
FILE_PATH_GLOB = "/files/data/pokemon_jpg/*.jpg" # replacing with local path

class ExternalInputIterator(object):
    def __init__(self, batch_size):
        self.files = glob.glob(FILE_PATH_GLOB)
        self.batch_size = batch_size
        shuffle(self.files)

    def __iter__(self):
        self.i = 0
        self.n = len(self.files)
        return self

    def __next__(self):
        batch = []
        for _ in range(self.batch_size):
            im = imageio.imread(self.files[self.i])
            im = np.asarray(im, dtype=np.float32)
            im = im / 255.0 # uncomment if line 43 is commented
            batch.append(im)
            self.i = (self.i + 1) % self.n
        return (batch,)

# creating and testing iterator
eii = ExternalInputIterator(BATCH_SIZE)
print(type(next(iter(eii))[0][0]))

# creating pipeline
pipe = Pipeline(batch_size=BATCH_SIZE, num_threads=2, device_id=0)
with pipe:
    images = fn.external_source(source=eii, num_outputs=1, device="cpu")
    #images = images / 255.0 # uncomment if line 30 is commented
    images = fn.resize(images, size=[256,256])
    pipe.set_outputs(images)
pipe.build()

# testing pipeline
pipe_out = pipe.run()
batch = pipe_out[0]
img = batch.at(0)
print(img.shape)

# trying to create DALIDataset
shapes = ((BATCH_SIZE, 256, 256, 3),)
dtypes = (tf.float32,)

with tf.device('/cpu:0'):
    dataset = dali_tf.experimental.DALIDatasetWithInputs(
        pipeline=pipe,
        batch_size=BATCH_SIZE,
        output_shapes=shapes,
        output_dtypes=dtypes,
        device_id=0)