Repository: pytorch/vision

Issue:
allow utils.save_image to write to a file object
Currently, the `torchvision.utils.save_image` function only allows to save to a file on disk.
https://github.com/pytorch/vision/blob/26c9630bc7add88685400e54457203816c06e750/torchvision/utils.py#L90-L105

It cannot write to file objects as that requires the [`format` parameter](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save) to be passed which `torchvision.utils.save_image` does not support as of yet.

Can you add this parameter to `torchvision.utils.save_image`?


Discussion:

Resolving patch:
diff --git a/torchvision/utils.py b/torchvision/utils.py
--- a/torchvision/utils.py
+++ b/torchvision/utils.py
@@ -87,13 +87,16 @@ def norm_range(t, range):
     return grid
 
 
-def save_image(tensor, filename, nrow=8, padding=2,
-               normalize=False, range=None, scale_each=False, pad_value=0):
+def save_image(tensor, fp, nrow=8, padding=2,
+               normalize=False, range=None, scale_each=False, pad_value=0, format=None):
     """Save a given Tensor into an image file.
 
     Args:
         tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
             saves the tensor as a grid of images by calling ``make_grid``.
+        fp - A filename(string) or file object
+        format(Optional):  If omitted, the format to use is determined from the filename extension.
+            If a file object was used instead of a filename, this parameter should always be used.
         **kwargs: Other arguments are documented in ``make_grid``.
     """
     from PIL import Image
@@ -102,4 +105,4 @@ def save_image(tensor, filename, nrow=8, padding=2,
     # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
     ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
     im = Image.fromarray(ndarr)
-    im.save(filename)
+    im.save(fp, format=format)


Test patch:
diff --git a/test/test_utils.py b/test/test_utils.py
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -3,6 +3,9 @@
 import torch
 import torchvision.utils as utils
 import unittest
+from io import BytesIO
+import torchvision.transforms.functional as F
+from PIL import Image
 
 
 class Tester(unittest.TestCase):
@@ -49,6 +52,26 @@ def test_save_image_single_pixel(self):
             utils.save_image(t, f.name)
             assert os.path.exists(f.name), 'The pixel image is not present after save'
 
+    def test_save_image_file_object(self):
+        with tempfile.NamedTemporaryFile(suffix='.png') as f:
+            t = torch.rand(2, 3, 64, 64)
+            utils.save_image(t, f.name)
+            img_orig = Image.open(f.name)
+            fp = BytesIO()
+            utils.save_image(t, fp, format='png')
+            img_bytes = Image.open(fp)
+            assert torch.equal(F.to_tensor(img_orig), F.to_tensor(img_bytes)), 'Image not stored in file object'
+
+    def test_save_image_single_pixel_file_object(self):
+        with tempfile.NamedTemporaryFile(suffix='.png') as f:
+            t = torch.rand(1, 3, 1, 1)
+            utils.save_image(t, f.name)
+            img_orig = Image.open(f.name)
+            fp = BytesIO()
+            utils.save_image(t, fp, format='png')
+            img_bytes = Image.open(fp)
+            assert torch.equal(F.to_tensor(img_orig), F.to_tensor(img_bytes)), 'Pixel Image not stored in file object'
+
 
 if __name__ == '__main__':
     unittest.main()