diff --git a/docs/source/io.rst b/docs/source/io.rst index e7aeedc0716..6c21e68cf59 100644 --- a/docs/source/io.rst +++ b/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 diff --git a/torchvision/io/__init__.py b/torchvision/io/__init__.py index d92c54b7df3..029dd311592 100644 --- a/torchvision/io/__init__.py +++ b/torchvision/io/__init__.py @@ -14,6 +14,15 @@ 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 @@ "_read_video_clip_from_memory", "_read_video_meta_data", "VideoMetaData", - "Timebase" + "Timebase", + 'read_image', + 'decode_image', + 'encode_jpeg', + 'write_jpeg', + 'encode_png', + 'write_png', ] diff --git a/torchvision/io/image.py b/torchvision/io/image.py index 48ebe8b11ad..398d682689e 100644 --- a/torchvision/io/image.py +++ b/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,12 +167,16 @@ 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 - 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. + 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: raise ValueError('Image quality should be a positive number ' @@ -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)