Update docstings and expose main image io functions in the vision docs (#2760)

* Expose image io functions in docs

* Expose image funcs in module

* Update docstring

* Update docstrings for all functions

* Update section header

docs/source/io.rst

@@ -4,7 +4,8 @@ torchvision.io
 .. currentmodule:: torchvision.io
 
 The :mod:`torchvision.io` package provides functions for performing IO
-operations. They are currently specific to reading and writing video.
+operations. They are currently specific to reading and writing video and
+images.
 
 Video
 -----
@@ -14,3 +15,19 @@ Video
 .. autofunction:: read_video_timestamps
 
 .. autofunction:: write_video
+
+
+Image
+-----
+
+.. autofunction:: read_image
+
+.. autofunction:: decode_image
+
+.. autofunction:: encode_jpeg
+
+.. autofunction:: write_jpeg
+
+.. autofunction:: encode_png
+
+.. autofunction:: write_png

torchvision/io/__init__.py

@@ -14,6 +14,15 @@ from .video import (
     read_video_timestamps,
     write_video,
 )
+from .image import (
+    read_image,
+    decode_image,
+    encode_jpeg,
+    write_jpeg,
+    encode_png,
+    write_png
+)
+
 
 __all__ = [
     "write_video",
@@ -29,5 +38,11 @@ __all__ = [
     "_read_video_clip_from_memory",
     "_read_video_meta_data",
     "VideoMetaData",
-    "Timebase"
+    "Timebase",
+    'read_image',
+    'decode_image',
+    'encode_jpeg',
+    'write_jpeg',
+    'encode_png',
+    'write_png',
 ]

torchvision/io/image.py

@@ -73,14 +73,20 @@ def encode_png(input: torch.Tensor, compression_level: int = 6) -> torch.Tensor:
     """
     Takes an input tensor in CHW layout and returns a buffer with the contents
     of its corresponding PNG file.
-    Arguments:
-        input (Tensor[channels, image_height, image_width]): int8 image tensor
-    of `c` channels, where `c` must 3 or 1.
-        compression_level (int): Compression factor for the resulting file, it
-    must be a number between 0 and 9. Default: 6
+
+    Parameters
+    ----------
+    input: Tensor[channels, image_height, image_width]
+        int8 image tensor of `c` channels, where `c` must 3 or 1.
+    compression_level: int
+        Compression factor for the resulting file, it must be a number
+        between 0 and 9. Default: 6
+
     Returns
-        output (Tensor[1]): A one dimensional int8 tensor that contains the raw
-    bytes of the PNG file.
+    -------
+    output: Tensor[1]
+        A one dimensional int8 tensor that contains the raw bytes of the
+        PNG file.
     """
     output = torch.ops.image.encode_png(input, compression_level)
     return output
@@ -90,12 +96,16 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
     """
     Takes an input tensor in CHW layout (or HW in the case of grayscale images)
     and saves it in a PNG file.
-    Arguments:
-        input (Tensor[channels, image_height, image_width]): int8 image tensor
-    of `c` channels, where `c` must be 1 or 3.
-        filename (str): Path to save the image.
-        compression_level (int): Compression factor for the resulting file, it
-    must be a number between 0 and 9. Default: 6
+
+    Parameters
+    ----------
+    input: Tensor[channels, image_height, image_width]
+        int8 image tensor of `c` channels, where `c` must be 1 or 3.
+    filename: str
+        Path to save the image.
+    compression_level: int
+        Compression factor for the resulting file, it must be a number
+        between 0 and 9. Default: 6
     """
     torch.ops.image.write_png(input, filename, compression_level)
 
@@ -131,14 +141,20 @@ def encode_jpeg(input: torch.Tensor, quality: int = 75) -> torch.Tensor:
     """
     Takes an input tensor in CHW layout and returns a buffer with the contents
     of its corresponding JPEG file.
-    Arguments:
-        input (Tensor[channels, image_height, image_width]): int8 image tensor
-    of `c` channels, where `c` must be 1 or 3.
-        quality (int): Quality of the resulting JPEG file, it must be a number
-    between 1 and 100. Default: 75
+
+    Parameters
+    ----------
+    input: Tensor[channels, image_height, image_width])
+        int8 image tensor of `c` channels, where `c` must be 1 or 3.
+    quality: int
+        Quality of the resulting JPEG file, it must be a number between
+        1 and 100. Default: 75
+
     Returns
-        output (Tensor[1]): A one dimensional int8 tensor that contains the raw
-    bytes of the JPEG file.
+    -------
+    output: Tensor[1]
+        A one dimensional int8 tensor that contains the raw bytes of the
+        JPEG file.
     """
     if quality < 1 or quality > 100:
         raise ValueError('Image quality should be a positive number '
@@ -151,11 +167,15 @@ def encode_jpeg(input: torch.Tensor, quality: int = 75) -> torch.Tensor:
 def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
     """
     Takes an input tensor in CHW layout and saves it in a JPEG file.
-    Arguments:
-        input (Tensor[channels, image_height, image_width]): int8 image tensor
-    of `c` channels, where `c` must be 1 or 3.
-        filename (str): Path to save the image.
-        quality (int): Quality of the resulting JPEG file, it must be a number
+
+    Parameters
+    ----------
+    input: Tensor[channels, image_height, image_width]
+        int8 image tensor of `c` channels, where `c` must be 1 or 3.
+    filename: str
+        Path to save the image.
+    quality: int
+        Quality of the resulting JPEG file, it must be a number
         between 1 and 100. Default: 75
     """
     if quality < 1 or quality > 100:
@@ -172,11 +192,15 @@ def decode_image(input: torch.Tensor) -> torch.Tensor:
 
     The values of the output tensor are uint8 between 0 and 255.
 
-    Arguments:
-        input (Tensor): a one dimensional uint8 tensor containing
-        the raw bytes of the PNG or JPEG image.
-    Returns:
-        output (Tensor[3, image_height, image_width])
+    Parameters
+    ----------
+    input: Tensor
+        a one dimensional uint8 tensor containing the raw bytes of the
+        PNG or JPEG image.
+
+    Returns
+    -------
+    output: Tensor[3, image_height, image_width]
     """
     output = torch.ops.image.decode_image(input)
     return output
@@ -186,10 +210,15 @@ def read_image(path: str) -> torch.Tensor:
     """
     Reads a JPEG or PNG image into a 3 dimensional RGB Tensor.
     The values of the output tensor are uint8 between 0 and 255.
-    Arguments:
-        path (str): path of the JPEG or PNG image.
-    Returns:
-        output (Tensor[3, image_height, image_width])
+
+    Parameters
+    ----------
+    path: str
+        path of the JPEG or PNG image.
+
+    Returns
+    -------
+    output: Tensor[3, image_height, image_width]
     """
     data = read_file(path)
     return decode_image(data)