Replace this
IMAGE_NAME = next(iter(dataset.images.keys()))
with
IMAGE_NAME = dataset.image_paths[0]
or if you want to iterate
IMAGE_NAME = next(iter(dataset.image_paths))
Hope this will fix the Attribute error