How resize dataset label in albumentations label to work with tensorflow image_dataset_from_directory function?



I am running the following code:
The error obtained is following:

File "C:\Users\Admin\PycharmProjects\BugsClassfications\",
line 45, in set_shapes *

ValueError: Shapes must be equal rank, but are 1 and 0

How correct function set_shape to work with image_dataset_from_directory?

Here is my code:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from functools import partial
from albumentations import (Compose, HorizontalFlip,Rotate)
def process_image(image, label, img_size):
  # cast and normalize image
  image = tf.image.convert_image_dtype(image, tf.float32)
  # apply simple augmentations
  image = tf.image.random_flip_left_right(image)
  image = tf.image.resize(image,[img_size, img_size])
  return image, label
transforms = Compose([
def aug_fn(image, img_size):
  data = {"image":image}
  aug_data = transforms(**data)
  aug_img = aug_data["image"]
  aug_img = tf.cast(aug_img/255.0, tf.float32)
  aug_img = tf.image.resize(aug_img, size=[img_size, img_size])
  return aug_img
def process_data(image, label, img_size):
  aug_img = tf.numpy_function(func=aug_fn, inp=[image, img_size], Tout=tf.float32)
  return aug_img, label
def set_shapes(img, label, img_shape=(128,128,3)):
  return img, label
def view_image(ds):
  image, label = next(iter(ds))  # extract 1 batch from the dataset
  image = image.numpy()
  label = label.numpy()
  fig = plt.figure(figsize=(22, 22))
  for i in range(20):
    ax = fig.add_subplot(4, 5, i + 1, xticks=[], yticks=[])
    ax.set_title(f"Label: {label[i]}")
train_dir = './dataset/train'
img_size = 128
data = tf.keras.utils.image_dataset_from_directory(train_dir, image_size=(img_size, img_size))
ds_alb =, img_size = 128), num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)
ds_alb =, num_parallel_calls=AUTOTUNE).batch(32)


If you change the shape of your labels, it should work:

def set_shapes(img, label, img_shape=(128,128,3)):
  return img, label

But you should ask yourself why you are even explicitly setting the shape of your data. Check this post.

Answered By – AloneTogether

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave A Reply

Your email address will not be published.

This website uses cookies to improve your experience. We'll assume you're ok with this, but you can opt-out if you wish. Accept Read More