Accept single integers for image_size in image_dataset_for_directory

keras-team · Aug 25, 2024 · 6adc2bf · 6adc2bf
1 parent 425ebc8
commit 6adc2bf
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/keras/src/utils/image_dataset_utils.py b/keras/src/utils/image_dataset_utils.py
@@ -197,6 +197,14 @@ def image_dataset_from_directory(
             f"Received: color_mode={color_mode}"
         )
 
+    if isinstance(image_size, int):
+        image_size = (image_size, image_size)
+    elif not isinstance(image_size, (list, tuple)) or not len(image_size) == 2:
+        raise ValueError(
+            "Invalid `image_size` value. Expected a tuple of 2 integers. "
+            f"Received: image_size={image_size}"
+        )
+
     interpolation = interpolation.lower()
     supported_interpolations = (
         "bilinear",