Support torch tensors

imgcat/imgcat.py

@@ -105,17 +105,21 @@ def to_content_buf(data):
         else:
             raise ValueError("Expected a 3D ndarray (RGB/RGBA image) or 2D (grayscale image), "
                              "but given shape: {}".format(im.shape))
+        return _get_bytes_from_numpy(im, mode)
 
-        try:
-            from PIL import Image
-        except ImportError as e:
-            raise ImportError(e.msg +
-                              "\nTo draw numpy arrays, we require Pillow. " +
-                              "(pip install Pillow)")       # TODO; reraise
-
-        with io.BytesIO() as buf:
-            Image.fromarray(im, mode=mode).save(buf, format='png')
-            return buf.getvalue()
+    elif 'torch' in sys.modules and isinstance(data, sys.modules['torch'].Tensor):
+        # numpy ndarray: convert to png
+        im = data
+        if im.shape[0] == 1:
+            mode = 'L'     # 8-bit pixels, grayscale
+            im = im.mul(255).byte().squeeze().numpy()
+        elif im.shape[0] == 3:
+            mode = None    # RGB/RGBA
+            im = im.mul(255).byte().permute(1, 2, 0).numpy()
+        else:
+            raise ValueError("Expected a 3D ndarray (RGB/RGBA image) or 2D (grayscale image), "
+                             "but given shape: {}".format(im.shape))
+        return _get_bytes_from_numpy(im, mode)
 
     elif 'PIL.Image' in sys.modules and isinstance(data, sys.modules['PIL.Image'].Image):
         # PIL/Pillow images
@@ -288,5 +292,18 @@ def main():
     return 0
 
 
+def _get_bytes_from_numpy(im, mode):
+    try:
+        from PIL import Image
+    except ImportError as e:
+        raise ImportError(e.msg +
+                          "\nTo draw numpy arrays, we require Pillow. " +
+                          "(pip install Pillow)")       # TODO; reraise
+
+    with io.BytesIO() as buf:
+        Image.fromarray(im, mode=mode).save(buf, format='png')
+        return buf.getvalue()
+
+
 if __name__ == '__main__':
     sys.exit(main())

imgcat/test_imgcat.py

@@ -2,6 +2,7 @@
 
 import unittest
 import numpy as np
+import torch
 import sys
 import os
 import io
@@ -41,6 +42,18 @@ def test_numpy(self):
         a[:, :, 0] = 255    # (255, 0, 0): red
         imgcat(a)
 
+    def test_torch(self):
+        # uint8, grayscale
+        a = torch.ones([1, 32, 32], dtype=torch.uint8)
+        imgcat(a)
+
+        a = torch.ones([1, 32, 32], dtype=torch.float32)
+        imgcat(a)
+
+        # uint8, color image
+        a = torch.ones([3, 32, 32], dtype=torch.uint8) * 0
+        imgcat(a)
+
     def test_matplotlib(self):
         # plt
         import matplotlib.pyplot as plt

setup.py

@@ -31,6 +31,7 @@ def read_version():
 tests_requires = [
     'pytest<5.0',
     'numpy',
+    'torch',
 ]
 if sys.version_info >= (3, 6):
     tests_requires += ['matplotlib>=3.1', 'Pillow']