Add pil_to_tensor to functionals

test/test_transforms.py

@@ -512,6 +512,50 @@ def test_accimage_to_tensor(self):
         self.assertEqual(expected_output.size(), output.size())
         self.assertTrue(np.allclose(output.numpy(), expected_output.numpy()))
 
+    def test_as_tensor(self):
+        test_channels = [1, 3, 4]
+        height, width = 4, 4
+        trans = transforms.AsTensor()
+
+        with self.assertRaises(TypeError):
+            trans(np.random.rand(1, height, width).tolist())
+
+        with self.assertRaises(ValueError):
+            trans(np.random.rand(height))
+            trans(np.random.rand(1, 1, height, width))
+
+        for channels in test_channels:
+            input_data = torch.ByteTensor(channels, height, width).random_(0, 255)
+            img = transforms.ToPILImage()(input_data)
+            output = trans(img)
+            self.assertTrue(np.allclose(input_data.numpy(), output.numpy()))
+
+            ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
+            output = trans(ndarray)
+            expected_output = ndarray.transpose((2, 0, 1))
+            self.assertTrue(np.allclose(output.numpy(), expected_output))
+
+            ndarray = np.random.rand(height, width, channels).astype(np.float32)
+            output = trans(ndarray)
+            expected_output = ndarray.transpose((2, 0, 1))
+            self.assertTrue(np.allclose(output.numpy(), expected_output))
+
+        # separate test for mode '1' PIL images
+        input_data = torch.ByteTensor(1, height, width).bernoulli_()
+        img = transforms.ToPILImage()(input_data.mul(255)).convert('1')
+        output = trans(img)
+        self.assertTrue(np.allclose(input_data.numpy(), output.numpy()))
+
+    @unittest.skipIf(accimage is None, 'accimage not available')
+    def test_accimage_as_tensor(self):
+        trans = transforms.AsTensor()
+
+        expected_output = trans(Image.open(GRACE_HOPPER).convert('RGB'))
+        output = trans(accimage.Image(GRACE_HOPPER))
+
+        self.assertEqual(expected_output.size(), output.size())
+        self.assertTrue(np.allclose(output.numpy(), expected_output.numpy()))
+
     @unittest.skipIf(accimage is None, 'accimage not available')
     def test_accimage_resize(self):
         trans = transforms.Compose([

torchvision/transforms/functional.py

@@ -82,6 +82,54 @@ def to_tensor(pic):
         return img
 
 
+def as_tensor(pic):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor of same type.
+
+    See ``AsTensor`` for more details.
+
+    Args:
+        pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+    Returns:
+        Tensor: Converted image.
+    """
+    if not(_is_pil_image(pic) or _is_numpy(pic)):
+        raise TypeError('pic should be PIL Image or ndarray. Got {}'.format(type(pic)))
+
+    if _is_numpy(pic) and not _is_numpy_image(pic):
+        raise ValueError('pic should be 2/3 dimensional. Got {} dimensions.'.format(pic.ndim))
+
+    if isinstance(pic, np.ndarray):
+        # handle numpy array
+        if pic.ndim == 2:
+            pic = pic[:, :, None]
+
+        img = torch.from_numpy(pic.transpose((2, 0, 1)))
+        return img
+
+    if accimage is not None and isinstance(pic, accimage.Image):
+        nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
+        pic.copyto(nppic)
+        return torch.from_numpy(nppic)
+
+    # handle PIL Image
+    if pic.mode == 'I':
+        img = torch.from_numpy(np.array(pic, np.int32, copy=False))
+    elif pic.mode == 'I;16':
+        img = torch.from_numpy(np.array(pic, np.int16, copy=False))
+    elif pic.mode == 'F':
+        img = torch.from_numpy(np.array(pic, np.float32, copy=False))
+    elif pic.mode == '1':
+        img = torch.from_numpy(np.array(pic, np.uint8, copy=False))
+    else:
+        img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
+
+    img = img.view(pic.size[1], pic.size[0], len(pic.getbands()))
+    # put it from HWC to CHW format
+    img = img.permute((2, 0, 1)).contiguous()
+    return img
+
+
 def to_pil_image(pic, mode=None):
     """Convert a tensor or an ndarray to PIL Image.
 

torchvision/transforms/transforms.py

@@ -15,7 +15,7 @@
 from . import functional as F
 
 
-__all__ = ["Compose", "ToTensor", "ToPILImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad",
+__all__ = ["Compose", "ToTensor", "AsTensor", "ToPILImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad",
            "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop", "RandomHorizontalFlip",
            "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation",
            "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale",
@@ -95,6 +95,28 @@ def __repr__(self):
         return self.__class__.__name__ + '()'
 
 
+class AsTensor(object):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor of the same type.
+
+    Converts a PIL Image or numpy.ndarray (H x W x C) to a torch.Tensor of shape (C x H x W)
+    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
+    or if the numpy.ndarray has dtype = np.uint8
+    """
+
+    def __call__(self, pic):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+        Returns:
+            Tensor: Converted image.
+        """
+        return F.as_tensor(pic)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
 class ToPILImage(object):
     """Convert a tensor or an ndarray to PIL Image.