In the previous article, we had a chance to see how one can scrape images from the web using Python. Apart from that, in one of the articles before that we could see how we can perform transfer learning with TensorFlow. In that article, we used famous Convolution Neural Networks on already prepared TensorFlow dataset. So, technically we are missing one step between scraping data from the web and training, right? How can we create TensorFlow dataset from images we just scraped from the web? In this article, we will do just that, prepare data and unify it under TensorFlow dataset.

For the purpose of this article and speeding up the process, we use one interesting source of images – Images of LEGO Bricks. Here you can find 16 different classes of LEGO bricks. Each brick is selected from Mecabricks and there are 400 different angles for each one of them. So, let’s imagine that we scraped these images from the web and now we want to create a TensorFlow dataset that is going to be used in learning process for neural network for classifying LEGO bricks. Watch your step 🙂


Once you download the images from the link above, you will notice that they are split into 16 directories, meaning there are 16 classes of LEGO bricks. If we were scraping these images, we would have to split them into these folders ourselves. This is important thing to do, since the all other steps depend on this. To sum it up, these all Lego Brick images are split into these folders:


In general, there are two ways we can achieve the goal. One is using Keras generator and the other is using pure TensorFlow core functionalities. No matter which approach do we choose, we need to import some libraries:

Apart from that, we need to load the path to the images and define classes. For that we use names of the folders in which images are located:

Here is the data which is in CLASSES variable:

array([‘11214 Bush 3M friction with Cross axle’, ‘18651 Cross Axle 2M with Snap friction’, ‘2357 Brick corner 1x2x2’, ‘3003 Brick 2×2’, ‘3004 Brick 1×2’, ‘3005 Brick 1×1’, ‘3022 Plate 2×2’, ‘3023 Plate 1×2’, ‘3024 Plate 1×1’, ‘3040 Roof Tile 1x2x45deg’, ‘3069 Flat Tile 1×2’, ‘32123 half Bush’, ‘3673 Peg 2M’, ‘3713 Bush for Cross Axle’, ‘3794 Plate 1X2 with 1 Knob’, ‘6632 Technic Lever 3M’], dtype='<U38′)

So, let’s first check out how we can create TensorFlow dataset with Keras using this information.

Keras Implementation

Creating dataset using Keras is pretty straight forward:

We are using ImageDataGenerator class from keras.preprocessing.image module. The only parameter we need in the constructor is rescale parameter. Using this we basically normalize all images. Once this object is created we call flow_from_firectory method. Here we pass on the path to the directory in which images are located and list of class names. We also pass on the information of the batch size, and the size to which all images will be resized.

This way we get 300×500 normalized images in the batches of 32 images. Here is how those images look like:

The next batch can be obtained like this:

TensorFlow Implementation

While Keras implemetation is quite easy, sometimes it’s performance can be bad. Meaning, that it can take some time while this is done. That is why we can do the same thing with pure TensorFlow. First thing that we need to do is get list of all image paths. That is done like this:

That way in list_dataset variable we have this info:

b’LEGO brick images\\train\\11214 Bush 3M friction with Cross axle\\201706171006-0003.png’ b’LEGO brick images\\train\\6632 Technic Lever 3M\\201706171606-0395.png’ b’LEGO brick images\\train\\3673 Peg 2M\\0362.png’ b’LEGO brick images\\train\\2357 Brick corner 1x2x2\\201706171206-0032.png’ b’LEGO brick images\\train\\3023 Plate 1×2\\0175.png’

Once that is done, we implement DataSetCreator class for the purpose of preparing images and the dataset. Here is what that looks like:

This class is initialized by batch size, image dimensions and the list of files. There are three private methods:

  • _get_class – Based on the path of the file, it retrieves the class of the image.
  • _load_image – Loads image from the defined path.
  • _load_labeled_data – Utilizes previous two functions, returns image data and it’s class (label).

However, majority of important stuff happens in load_process method. Let’s take a closer look:

In this function, we utilize map function and for each image file path that we previously loaded we call _load_labeled_data method. This in turn loads all images and it’s classes into self.loaded_dataset. Once this is done, we can cache and shuffle dataset. After that we create batches. Additional cool thing that we do is call the prefatch method on the dataset. This method let’s dataset work in the background. Basically, during the training process, dataset performs lazy loading of the images from the disk. This won’t slow down the training process. This is why this implementation is so cool.

Finally, we can create an object of the DataSetCreator class and use get_batch method to get the data:

The result is the same as with Keras implementation:


In this article, we created TensorFlow dataset using downloaded images. This dataset now can be used for training some neural networks or different classification algorithms.

Thank you for reading!

Read more posts from the author at Rubik’s Code.