Probability parameter in RandomHorizontalFlip, RandomHorizontalFlip (#417)

* Set probability as configuration parameter in RandomHorizontalFlip and RandomHorizontalFlip (#414)

* Fix documentation

Author: vfdev-5

test/test_transforms.py

@@ -475,6 +475,17 @@ def test_random_vertical_flip(self):
         random.setstate(random_state)
         assert p_value > 0.0001
 
+        num_samples = 250
+        num_vertical = 0
+        for _ in range(num_samples):
+            out = transforms.RandomVerticalFlip(p=0.7)(img)
+            if out == vimg:
+                num_vertical += 1
+
+        p_value = stats.binom_test(num_vertical, num_samples, p=0.7)
+        random.setstate(random_state)
+        assert p_value > 0.0001
+
         # Checking if RandomVerticalFlip can be printed as string
         transforms.RandomVerticalFlip().__repr__()
 
@@ -496,6 +507,17 @@ def test_random_horizontal_flip(self):
         random.setstate(random_state)
         assert p_value > 0.0001
 
+        num_samples = 250
+        num_horizontal = 0
+        for _ in range(num_samples):
+            out = transforms.RandomHorizontalFlip(p=0.7)(img)
+            if out == himg:
+                num_horizontal += 1
+
+        p_value = stats.binom_test(num_horizontal, num_samples, p=0.7)
+        random.setstate(random_state)
+        assert p_value > 0.0001
+
         # Checking if RandomHorizontalFlip can be printed as string
         transforms.RandomHorizontalFlip().__repr__()
 

torchvision/transforms/transforms.py

@@ -321,7 +321,14 @@ def __repr__(self):
 
 
 class RandomHorizontalFlip(object):
-    """Horizontally flip the given PIL Image randomly with a probability of 0.5."""
+    """Horizontally flip the given PIL Image randomly with a given probability.
+
+    Args:
+        p (float): probability of the image being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
 
     def __call__(self, img):
         """
@@ -331,16 +338,23 @@ def __call__(self, img):
         Returns:
             PIL Image: Randomly flipped image.
         """
-        if random.random() < 0.5:
+        if random.random() < self.p:
             return F.hflip(img)
         return img
 
     def __repr__(self):
-        return self.__class__.__name__ + '()'
+        return self.__class__.__name__ + '(p={})'.format(self.p)
 
 
 class RandomVerticalFlip(object):
-    """Vertically flip the given PIL Image randomly with a probability of 0.5."""
+    """Vertically flip the given PIL Image randomly with a given probability.
+
+    Args:
+        p (float): probability of the image being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
 
     def __call__(self, img):
         """
@@ -350,12 +364,12 @@ def __call__(self, img):
         Returns:
             PIL Image: Randomly flipped image.
         """
-        if random.random() < 0.5:
+        if random.random() < self.p:
             return F.vflip(img)
         return img
 
     def __repr__(self):
-        return self.__class__.__name__ + '()'
+        return self.__class__.__name__ + '(p={})'.format(self.p)
 
 
 class RandomResizedCrop(object):