Make crop scriptable

test/test_functional_tensor.py

@@ -1,6 +1,9 @@
+import torch
+import torchvision.transforms as transforms
 import torchvision.transforms.functional_tensor as F_t
+import torchvision.transforms.functional as F
 import unittest
-import torch
+import random
 
 
 class Tester(unittest.TestCase):
@@ -19,6 +22,21 @@ def test_hflip(self):
         self.assertEqual(hflipped_img.shape, img_tensor.shape)
         self.assertTrue(torch.equal(img_tensor, hflipped_img_again))
 
+    def test_crop(self):
+        img_tensor = torch.FloatTensor(3, 16, 16).uniform_(0, 1)
+        top = random.randint(0, 15)
+        left = random.randint(0, 15)
+        height = random.randint(1, 16-top)
+        width = random.randint(1, 16-left)
+        img_cropped = F_t.crop(img_tensor, top, left, height, width)
+        img_PIL = transforms.ToPILImage()(img_tensor)
+        img_PIL_cropped = F.crop(img_PIL, top, left, height, width)
+        img_cropped_GT = transforms.ToTensor()(img_PIL_cropped)
+
+        max_diff = (img_cropped_GT - img_cropped).abs().max()
+
+        assert max_diff < 5e-3, "Functional crop not working"
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/transforms/functional_tensor.py

@@ -12,7 +12,7 @@ def vflip(img_tensor):
         Tensor:  Vertically flipped image Tensor.
     """
     if not F._is_tensor_image(img_tensor):
-        raise TypeError('tensor is not a torch image.')
+        raise TypeError('Input image is not a tensor.')
 
     return img_tensor.flip(-2)
 
@@ -28,6 +28,23 @@ def hflip(img_tensor):
     """
 
     if not F._is_tensor_image(img_tensor):
-        raise TypeError('tensor is not a torch image.')
+        raise TypeError('Input image is not a tensor.')
 
     return img_tensor.flip(-1)
+
+
+def crop(img, top, left, height, width):
+    """Crop the given Image Tensor.
+    Args:
+        img (Tensor): Image to be cropped in the form [C, H, W]. (0,0) denotes the top left corner of the image.
+        top (int): Vertical component of the top left corner of the crop box.
+        left (int): Horizontal component of the top left corner of the crop box.
+        height (int): Height of the crop box.
+        width (int): Width of the crop box.
+    Returns:
+        Tensor: Cropped image.
+    """
+    if not F._is_tensor_image(img):
+        raise TypeError('Input image is not a tensor.')
+
+    return img[..., top:top + height, left:left + width]