[AIR] Add experimental read_images

doc/source/data/api/input_output.rst

@@ -48,6 +48,11 @@ Text
 
 .. autofunction:: ray.data.read_text
 
+Images (experimental)
+---------------------
+
+.. autofunction:: ray.data.read_images
+
 Binary
 ------
 
@@ -214,4 +219,4 @@ MetadataProvider API
     :members:
 
 .. autoclass:: ray.data.datasource.FastFileMetadataProvider
-    :members:
+    :members:

doc/source/data/creating-datasets.rst

@@ -162,6 +162,30 @@ Supported File Formats
 
   See the API docs for :func:`read_text() <ray.data.read_text>`.
 
+.. tabbed:: Images (experimental)
+
+  If  your directory structure is:
+
+  .. code-block::
+
+    root/dog/xxx.png
+    root/dog/xxy.png
+    root/dog/[...]/xxz.png
+
+    root/cat/123.png
+    root/cat/nsdf3.png
+    root/cat/[...]/asd932_.png
+
+  Then call :func:`~ray.data.read_images` to load your images into a ``Dataset``.
+
+  .. literalinclude:: ./doc_code/creating_datasets.py
+    :language: python
+    :start-after: __read_images_begin__
+    :end-before: __read_images_end__
+
+  For more information on working with tensors, see our
+  :ref:`tensor data guide <datasets_tensor_support>`
+
 .. tabbed:: Binary
 
   Read binary files into a ``Dataset``. Each binary file will be treated as a single row
@@ -518,19 +542,6 @@ converts it into a Ray Dataset directly.
     ray_datasets["train"].take(2)
     # [{'text': ''}, {'text': ' = Valkyria Chronicles III = \n'}]
 
-.. _datasets_from_images:
-
--------------------------------
-From Image Files (experimental)
--------------------------------
-
-Load image data stored as individual files using :py:class:`~ray.data.datasource.ImageFolderDatasource`:
-
-.. literalinclude:: ./doc_code/tensor.py
-    :language: python
-    :start-after: __create_images_begin__
-    :end-before: __create_images_end__
-
 .. _datasets_custom_datasource:
 
 ------------------

doc/source/data/dataset-tensor-support.rst

@@ -89,7 +89,7 @@ This section shows how to create single and multi-column Tensor datasets.
 
 .. tabbed:: Images (experimental)
 
-  Load image data stored as individual files using :py:class:`~ray.data.datasource.ImageFolderDatasource`:
+  Load image data stored as individual files using :func:`~ray.data.read_images`:
 
   **Image and label columns**:
 

doc/source/data/dataset.rst

@@ -110,7 +110,7 @@ Advanced users can refer directly to the Ray Datasets :ref:`API reference <data-
         :text: Start Using Ray Datasets
         :classes: btn-outline-info btn-block
     ---
-    
+
     **Examples**
     ^^^
 
@@ -200,6 +200,9 @@ Supported Input Formats
    * - Text Files
      - :func:`ray.data.read_text()`
      - ✅
+   * - Image Files (experimental)
+     - :func:`ray.data.read_images()`
+     - 🚧
    * - Binary Files
      - :func:`ray.data.read_binary_files()`
      - ✅

doc/source/data/doc_code/creating_datasets.py

@@ -149,6 +149,33 @@
 # __from_numpy_end__
 # fmt: on
 
+# fmt: off
+# __read_images_begin__
+ds = ray.data.read_images(root="example://image-folders/simple", size=(128, 128))
+# -> Dataset(num_blocks=3, num_rows=3,
+#            schema={image: TensorDtype(shape=(128, 128, 3), dtype=uint8),
+#                    label: object})
+
+ds.take(1)
+# -> [{'image':
+#         array([[[ 92,  71,  57],
+#                 [107,  87,  72],
+#                 ...,
+#                 [141, 161, 185],
+#                 [139, 158, 184]],
+#
+#                ...,
+#
+#                [[135, 135, 109],
+#                 [135, 135, 108],
+#                 ...,
+#                 [167, 150,  89],
+#                 [165, 146,  90]]], dtype=uint8),
+#      'label': 'cat',
+#     }]
+# __read_images_end__
+# fmt: on
+
 # fmt: off
 # __from_numpy_mult_begin__
 import numpy as np

doc/source/data/doc_code/tensor.py

@@ -194,10 +194,7 @@ def cast_udf(block: pa.Table) -> pa.Table:
 ds.fully_executed()
 
 # __create_images_begin__
-from ray.data.datasource import ImageFolderDatasource
-
-ds = ray.data.read_datasource(
-    ImageFolderDatasource(), root="example://image-folders/simple", size=(128, 128))
+ds = ray.data.read_images(root="example://image-folders/simple", size=(128, 128))
 # -> Dataset(num_blocks=3, num_rows=3,
 #            schema={image: TensorDtype(shape=(128, 128, 3), dtype=uint8),
 #                    label: object})

doc/source/ray-air/examples/torch_image_batch_pretrained.py

@@ -7,7 +7,6 @@
 from ray.train.torch import TorchCheckpoint, TorchPredictor
 from ray.train.batch_predictor import BatchPredictor
 from ray.data.preprocessors import BatchMapper
-from ray.data.datasource import ImageFolderDatasource
 
 
 def preprocess(df: pd.DataFrame) -> pd.DataFrame:
@@ -29,9 +28,7 @@ def preprocess(df: pd.DataFrame) -> pd.DataFrame:
 
 data_url = "s3://anonymous@air-example-data-2/1G-image-data-synthetic-raw"
 print(f"Running GPU batch prediction with 1GB data from {data_url}")
-dataset = ray.data.read_datasource(
-    ImageFolderDatasource(), root=data_url, size=(256, 256)
-)
+dataset = ray.data.read_images(root=data_url, size=(256, 256))
 
 model = resnet18(pretrained=True)
 

python/ray/data/__init__.py

@@ -29,6 +29,7 @@
     read_binary_files,
     read_csv,
     read_datasource,
+    read_images,
     read_json,
     read_numpy,
     read_parquet,
@@ -70,6 +71,7 @@
     "read_binary_files",
     "read_csv",
     "read_datasource",
+    "read_images",
     "read_json",
     "read_numpy",
     "read_parquet",

python/ray/data/datasource/image_folder_datasource.py

@@ -40,85 +40,15 @@
 
 @DeveloperAPI
 class ImageFolderDatasource(BinaryDatasource):
-    """A datasource that lets you read datasets like `ImageNet <https://www.image-net.org/>`_.
-
-    This datasource works with any dataset where images are arranged in this way:
-
-    .. code-block::
-
-        root/dog/xxx.png
-        root/dog/xxy.png
-        root/dog/[...]/xxz.png
-
-        root/cat/123.png
-        root/cat/nsdf3.png
-        root/cat/[...]/asd932_.png
-
-    Datasets read with this datasource contain two columns: ``'image'`` and ``'label'``.
-
-    * The ``'image'`` column is of type
-      :py:class:`~ray.air.util.tensor_extensions.pandas.TensorDtype`. The shape of the
-      tensors are :math:`(H, W)` if the images are grayscale and :math:`(H, W, C)`
-      otherwise.
-    * The ``'label'`` column contains strings representing class names (e.g., 'cat').
-
-    Examples:
-        >>> import ray
-        >>> from ray.data.datasource import ImageFolderDatasource
-        >>> ds = ray.data.read_datasource(  # doctest: +SKIP
-        ...     ImageFolderDatasource(),
-        ...     root="/data/imagenet/train",
-        ...     size=(224, 224)
-        ... )
-        >>> sample = ds.take(1)[0]  # doctest: +SKIP
-        >>> sample["image"].to_numpy().shape  # doctest: +SKIP
-        (224, 224, 3)
-        >>> sample["label"]  # doctest: +SKIP
-        'n01443537'
-
-        To convert class labels to integer-valued targets, use
-        :py:class:`~ray.data.preprocessors.OrdinalEncoder`.
-
-        >>> import ray
-        >>> from ray.data.preprocessors import OrdinalEncoder
-        >>> ds = ray.data.read_datasource(  # doctest: +SKIP
-        ...     ImageFolderDatasource(),
-        ...     root="/data/imagenet/train",
-        ...     size=(224, 224)
-        ... )
-        >>> oe = OrdinalEncoder(columns=["label"])  # doctest: +SKIP
-        >>> ds = oe.fit_transform(ds)  # doctest: +SKIP
-        >>> sample = ds.take(1)[0]  # doctest: +SKIP
-        >>> sample["label"]  # doctest: +SKIP
-        71
-    """  # noqa: E501
+    """A datasource that lets you read datasets like ImageNet."""
 
     def create_reader(
         self,
         root: str,
         size: Optional[Tuple[int, int]] = None,
         mode: Optional[str] = None,
     ) -> "Reader[T]":
-        """Return a :py:class:`~ray.data.datasource.Reader` that reads images.
-
-        .. warning::
-            If your dataset contains images of varying sizes and you don't specify
-            ``size``, this datasource will error. To prevent errors, specify ``size``
-            or :ref:`disable tensor extension casting <disable_tensor_extension_casting>`.
-
-        Args:
-            root: Path to the dataset root.
-            size: The desired height and width of loaded images. If unspecified, images
-                retain their original shape.
-            mode: A `Pillow mode <https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes>`_
-                describing the desired type and depth of pixels. If unspecified, image
-                modes are inferred by
-                `Pillow <https://pillow.readthedocs.io/en/stable/index.html>`_.
-
-        Raises:
-            ValueError: if ``size`` contains non-positive numbers.
-            ValueError: if ``mode`` is unsupported.
-        """  # noqa: E501
+        """Return a :py:class:`~ray.data.datasource.Reader` that reads images."""
         if size is not None and len(size) != 2:
             raise ValueError(
                 "Expected `size` to contain 2 integers for height and width, "

python/ray/data/read_api.py

@@ -27,6 +27,7 @@
     DefaultFileMetadataProvider,
     DefaultParquetMetadataProvider,
     FastFileMetadataProvider,
+    ImageFolderDatasource,
     JSONDatasource,
     NumpyDatasource,
     ParquetBaseDatasource,
@@ -377,6 +378,75 @@ def read_parquet(
     )
 
 
+@PublicAPI(stability="alpha")
+def read_images(
+    root: str, size: Optional[Tuple[int, int]] = None, mode: Optional[str] = None
+):
+    """Read datasets like `ImageNet <https://www.image-net.org/>`_.
+
+    This function works with any directory where images are arranged in this way:
+
+    .. code-block::
+
+        root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/[...]/xxz.png
+
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/[...]/asd932_.png
+
+    Datasets read with this function contain two columns: ``'image'`` and ``'label'``.
+
+    * The ``'image'`` column is of type
+      :py:class:`~ray.air.util.tensor_extensions.pandas.TensorDtype`. The shape of the
+      tensors are :math:`(H, W)` if the images are grayscale and :math:`(H, W, C)`
+      otherwise.
+    * The ``'label'`` column contains strings representing class names (e.g., 'cat').
+
+    .. warning::
+        If your dataset contains images of varying sizes and you don't specify
+        ``size``, this function will error. To prevent errors, specify ``size``
+        or :ref:`disable tensor extension casting <disable_tensor_extension_casting>`.
+
+    Examples:
+        >>> import ray
+        >>> ds = ray.data.read_images("/data/imagenet/train", size=(224, 224))
+        >>> sample = ds.take(1)[0]  # doctest: +SKIP
+        >>> sample["image"].to_numpy().shape  # doctest: +SKIP
+        (224, 224, 3)
+        >>> sample["label"]  # doctest: +SKIP
+        'n01443537'
+
+        To convert class labels to integer-valued targets, use
+        :class:`~ray.data.preprocessors.OrdinalEncoder`.
+
+        >>> from ray.data.preprocessors import OrdinalEncoder
+        >>> oe = OrdinalEncoder(columns=["label"])  # doctest: +SKIP
+        >>> ds = oe.fit_transform(ds)  # doctest: +SKIP
+        >>> sample = ds.take(1)[0]  # doctest: +SKIP
+        >>> sample["label"]  # doctest: +SKIP
+        71
+
+    Args:
+        root: Path to the dataset root.
+        size: The desired height and width of loaded images. If unspecified, images
+            retain their original shape.
+        mode: A `Pillow mode <https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes>`_
+            describing the desired type and depth of pixels. If unspecified, image
+            modes are inferred by
+            `Pillow <https://pillow.readthedocs.io/en/stable/index.html>`_.
+
+    Returns:
+        A :class:`~ray.data.Dataset` containing image and label columns.
+
+    Raises:
+        ValueError: if ``size`` contains non-positive numbers.
+        ValueError: if ``mode`` is unsupported.
+    """  # noqa: E501
+    return read_datasource(ImageFolderDatasource(), root=root, size=size, mode=mode)
+
+
 @PublicAPI
 def read_parquet_bulk(
     paths: Union[str, List[str]],