updated ToTensor to support more types

Author: bodokaiser

test/test_transforms.py

@@ -151,6 +151,30 @@ def test_tensor_to_pil_image(self):
         expected_output = img_data.mul(255).int().float().div(255)
         assert np.allclose(expected_output[0].numpy(), to_tensor(l).numpy())
 
+    def test_tensor_gray_to_pil_image(self):
+        trans = transforms.ToPILImage()
+        to_tensor = transforms.ToTensor()
+
+        img_data_byte = torch.ByteTensor(1, 4, 4).random_(0, 255)
+        img_data_short = torch.ShortTensor(1, 4, 4).random_()
+        img_data_int = torch.IntTensor(1, 4, 4).random_()
+        img_data_float = torch.FloatTensor(1, 4, 4).uniform_()
+
+        img_byte = trans(img_data_byte)
+        img_short = trans(img_data_short)
+        img_int = trans(img_data_int)
+        img_float = trans(img_data_float)
+        assert img_byte.mode == 'L'
+        assert img_short.mode == 'I;16'
+        assert img_int.mode == 'I'
+        #assert img_float.mode == 'F'
+
+        assert np.allclose(img_data_short.numpy(), to_tensor(img_short).numpy())
+        assert np.allclose(img_data_int.numpy(), to_tensor(img_int).numpy())
+        # would cause breaking changes as ToTensor converts to range [0, 1]
+        #assert np.allclose(img_data_byte.numpy(), to_tensor(img_byte).numpy())
+        #assert np.allclose(img_data_float.numpy(), to_tensor(img_float).numpy())
+
     def test_ndarray_to_pil_image(self):
         trans = transforms.ToPILImage()
         img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()

torchvision/transforms.py

@@ -39,19 +39,30 @@ def __call__(self, pic):
         if isinstance(pic, np.ndarray):
             # handle numpy array
             img = torch.from_numpy(pic.transpose((2, 0, 1)))
+            # backard compability
+            return img.float().div(255)
+        # handle PIL Image
+        if pic.mode == 'I':
+            img = torch.from_numpy(np.array(pic, np.int32))
+        elif pic.mode == 'I;16':
+            img = torch.from_numpy(np.array(pic, np.int16))
         else:
-            # handle PIL Image
             img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
-            # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
-            if pic.mode == 'YCbCr':
-                nchannel = 3
-            else:
-                nchannel = len(pic.mode)
-            img = img.view(pic.size[1], pic.size[0], nchannel)
-            # put it from HWC to CHW format
-            # yikes, this transpose takes 80% of the loading time/CPU
-            img = img.transpose(0, 1).transpose(0, 2).contiguous()
-        return img.float().div(255)
+        # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
+        if pic.mode == 'YCbCr':
+            nchannel = 3
+        elif pic.mode == 'I;16':
+            nchannel = 1
+        else:
+            nchannel = len(pic.mode)
+        img = img.view(pic.size[1], pic.size[0], nchannel)
+        # put it from HWC to CHW format
+        # yikes, this transpose takes 80% of the loading time/CPU
+        img = img.transpose(0, 1).transpose(0, 2).contiguous()
+        if isinstance(img, torch.ByteTensor):
+            return img.float().div(255)
+        else:
+            return img
 
 
 class ToPILImage(object):
@@ -67,7 +78,6 @@ def __call__(self, pic):
         if torch.is_tensor(pic):
             npimg = np.transpose(pic.numpy(), (1, 2, 0))
         assert isinstance(npimg, np.ndarray), 'pic should be Tensor or ndarray'
-
         if npimg.shape[2] == 1:
             npimg = npimg[:, :, 0]
 
@@ -83,7 +93,6 @@ def __call__(self, pic):
             if npimg.dtype == np.uint8:
                 mode = 'RGB'
         assert mode is not None, '{} is not supported'.format(npimg.dtype)
-
         return Image.fromarray(npimg, mode=mode)