How to load images from URL with a TensorFlow 2 dataset

Viewed 2313

I want to use a TensorFlow 2 dataset object to feed images to a CNN. My images are located on AWS S3, but I will use images from Wikipedia in my example (the problem is the same).

image_urls = [
    'https://upload.wikimedia.org/wikipedia/commons/6/60/Matterhorn_from_Domh%C3%BCtte_-_2.jpg',
    'https://upload.wikimedia.org/wikipedia/commons/6/6e/Matterhorn_from_Klein_Matterhorn.jpg',
]
dataset = tf.data.Dataset.from_tensor_slices(image_urls)

def read_image_from_url(url):
    img_array = None
    with urlopen(url) as request:
        img_array = np.asarray(bytearray(request.read()), dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  #as RGB image (cv2 is BGR by default)

When I test my function with one element of the dataset, it works:

url = next(iter(dataset)).numpy().decode('utf-8')
img = read_image_from_url(url)
plt.imshow(img)

But when I map my function to the dataset to create a new dataset serving the images, it fails:

dataset_images = dataset.map(lambda x: read_image_from_url(x.numpy().decode('utf-8')))

AttributeError: in converted code:

    <ipython-input-6-e8eb89833196>:2 None  *
        map_func=lambda x: read_image_from_url(x.numpy().decode('utf-8')),

    AttributeError: 'Tensor' object has no attribute 'numpy'

Obviously, the dataset provides a different dtype when iterated with next or with map. Any idea how I can fix this?

1 Answers

Well this was way harder than it needs to be:

import tensorflow as tf
import numpy as np 
import cv2
from urllib.request import urlopen
import matplotlib.pyplot as plt
image_urls = [
    'https://upload.wikimedia.org/wikipedia/commons/6/60/Matterhorn_from_Domh%C3%BCtte_-_2.jpg',
    'https://upload.wikimedia.org/wikipedia/commons/6/6e/Matterhorn_from_Klein_Matterhorn.jpg',
]
dataset = tf.data.Dataset.from_tensor_slices(image_urls)

def get(url):
    with urlopen(str(url.numpy().decode("utf-8"))) as request:
        img_array = np.asarray(bytearray(request.read()), dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

def read_image_from_url(url):
    return tf.py_function(get, [url], tf.uint8)


dataset_images = dataset.map(lambda x: read_image_from_url(x))

for d in dataset_images:
  print(d)

Why did the first one work and then failed in the tf.Dataset? Well tf.Dataset is defined in graph mode not in eager mode like the first one. Graph mode is faster, and tf.Dataset is optimized for speed, so it makes sense. You can't do .numpy(), in graph mode, as everything should be defined in tensorflow ops. py_func wraps a python function in an tf.Operation that is executed in eager mode, which is just what we needed.

Note: I tried tf.keras.utils.get_file(), but I ran into similar issues as you are describing here. Hope this helps!

Related