diff --git a/test/test_image.py b/test/test_image.py index 53714342d46..3d9d612b5f3 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -532,5 +532,21 @@ def test_write_jpeg(img_path, tmpdir, scripted): assert_equal(torch_bytes, pil_bytes) +def test_pathlib_support(tmpdir): + # Just make sure pathlib.Path is supported where relevant + + jpeg_path = Path(next(get_images(ENCODE_JPEG, ".jpg"))) + + read_file(jpeg_path) + read_image(jpeg_path) + + write_path = Path(tmpdir) / "whatever" + img = torch.randint(0, 10, size=(3, 4, 4), dtype=torch.uint8) + + write_file(write_path, data=img.flatten()) + write_jpeg(img, write_path) + write_png(img, write_path) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/torchvision/io/image.py b/torchvision/io/image.py index ff5e8c0e4d1..8d3b294b32e 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -42,14 +42,14 @@ def read_file(path: str) -> torch.Tensor: with one dimension. Args: - path (str): the path to the file to be read + path (str or ``pathlib.Path``): the path to the file to be read Returns: data (Tensor) """ if not torch.jit.is_scripting() and not torch.jit.is_tracing(): _log_api_usage_once(read_file) - data = torch.ops.image.read_file(path) + data = torch.ops.image.read_file(str(path)) return data @@ -59,12 +59,12 @@ def write_file(filename: str, data: torch.Tensor) -> None: file. Args: - filename (str): the path to the file to be written + filename (str or ``pathlib.Path``): the path to the file to be written data (Tensor): the contents to be written to the output file """ if not torch.jit.is_scripting() and not torch.jit.is_tracing(): _log_api_usage_once(write_file) - torch.ops.image.write_file(filename, data) + torch.ops.image.write_file(str(filename), data) def decode_png( @@ -123,7 +123,7 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6): Args: 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. + filename (str or ``pathlib.Path``): 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 """ @@ -211,7 +211,7 @@ def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75): Args: 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. + filename (str or ``pathlib.Path``): Path to save the image. quality (int): Quality of the resulting JPEG file, it must be a number between 1 and 100. Default: 75 """ @@ -259,7 +259,7 @@ def read_image( The values of the output tensor are uint8 in [0, 255]. Args: - path (str): path of the JPEG or PNG image. + path (str or ``pathlib.Path``): path of the JPEG or PNG image. mode (ImageReadMode): the read mode used for optionally converting the image. Default: ``ImageReadMode.UNCHANGED``. See ``ImageReadMode`` class for more information on various