import numpy as np
from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
import nvidia.dali.types as types


class ExternalInputCallable:
    def __init__(self):
        pass

    def __call__(self, sample_info):
        sample_idx = sample_info.idx_in_epoch
        if sample_idx >= 1000:
            # Indicate end of the epoch
            raise StopIteration()

        with open('./n03287733_3194.JPEG', 'rb') as fr:
            encoded_img = np.frombuffer(fr.read(), dtype=np.uint8)
        return encoded_img, np.array(sample_idx)


@pipeline_def(
    batch_size=1, num_threads=1, device_id=0, py_num_workers=1
)
def loader_pipeline():
    jpegs, inds = fn.external_source(
        source=ExternalInputCallable(),
        num_outputs=2,
        batch=False,
        parallel=True,
        dtype=[types.UINT8, types.INT64],
    )
    decode = fn.decoders.image(jpegs, device="mixed")
    return decode, inds



pipe = loader_pipeline()
pipe.start_py_workers()
pipe.build()


for i in range(10):
    with warnings.catch_warnings():
        warnings.simplefilter('error')
        try:
            one_out = pipe.run()
            img = np.array(one_out[0][0].as_cpu())
            ind = np.array(one_out[1][0])
        except Exception as e:
            print(i, ': ', str(e))