Merge pull request #121 from lewisbelcher/fix-torch-transform-bug

Update torch_transform and tests

Augmentor/Pipeline.py

@@ -627,11 +627,9 @@ def torch_transform(self):
         """
         def _transform(image):
             for operation in self.operations:
-                r = round(random.uniform(0, 1), 1)
-                if r <= operation.probability:
-                    image = [image]
-                    image = operation.perform_operation(image)
-
+                r = random.uniform(0, 1)
+                if r < operation.probability:
+                    image = operation.perform_operation([image])[0]
             return image
 
         return _transform

tests/test_torch_transform.py

@@ -14,13 +14,21 @@ def test_torch_transform():
     red[..., 0] = 255
     red = Image.fromarray(red)
 
-    g = Augmentor.Operations.Greyscale(1)
-
     p = Augmentor.Pipeline()
-    p.greyscale(1)
+
+    # include multiple transforms to test integration
+    p.greyscale(probability=1)
+    p.zoom(probability=1, min_factor=1.0, max_factor=1.0)
+    p.rotate_random_90(probability=1)
+
     transforms = torchvision.transforms.Compose([
         p.torch_transform()
     ])
 
     assert red != transforms(red)
-    assert g.perform_operation([red]) == transforms(red)
+
+    # assert that all operations were correctly applied
+    result = red
+    for op in p.operations:
+        result = op.perform_operation([result])[0]
+    assert transforms(red) == result