update random apply to work also with targets

mindee · Sep 27, 2023 · 3443def · 3443def
1 parent 56395ba
commit 3443def
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 8 deletions.
diff --git a/doctr/transforms/modules/base.py b/doctr/transforms/modules/base.py
@@ -5,7 +5,7 @@
 
 import math
 import random
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
@@ -206,10 +206,10 @@ def __init__(self, transform: Callable[[Any], Any], p: float = 0.5) -> None:
     def extra_repr(self) -> str:
         return f"transform={self.transform}, p={self.p}"
 
-    def __call__(self, img: Any) -> Any:
+    def __call__(self, img: Any, target: Optional[np.ndarray] = None) -> Union[Any, Tuple[Any, np.ndarray]]:
         if random.random() < self.p:
-            return self.transform(img)
-        return img
+            return self.transform(img) if target is None else self.transform(img, target)
+        return img if target is None else (img, target)
 
 
 class RandomRotate(NestedObject):

diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
@@ -193,7 +193,7 @@ def main(args):
             + (
                 [
                     T.Resize(args.input_size, preserve_aspect_ratio=True),  # This does not pad
-                    T.RandomRotate(90, expand=True),
+                    T.RandomApply(T.RandomRotate(90, expand=True), 0.5),
                     T.Resize((args.input_size, args.input_size), preserve_aspect_ratio=True, symmetric_pad=True),
                 ]
                 if args.rotation and not args.eval_straight
@@ -286,7 +286,7 @@ def main(args):
             + (
                 [
                     T.Resize(args.input_size, preserve_aspect_ratio=True),
-                    T.RandomRotate(90, expand=True),
+                    T.RandomApply(T.RandomRotate(90, expand=True), 0.5),
                     T.Resize((args.input_size, args.input_size), preserve_aspect_ratio=True, symmetric_pad=True),
                 ]
                 if args.rotation

diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
@@ -155,7 +155,7 @@ def main(args):
             + (
                 [
                     T.Resize(args.input_size, preserve_aspect_ratio=True),  # This does not pad
-                    T.RandomRotate(90, expand=True),
+                    T.RandomApply(T.RandomRotate(90, expand=True), 0.5),
                     T.Resize((args.input_size, args.input_size), preserve_aspect_ratio=True, symmetric_pad=True),
                 ]
                 if args.rotation and not args.eval_straight
@@ -240,7 +240,7 @@ def main(args):
             + (
                 [
                     T.Resize(args.input_size, preserve_aspect_ratio=True),  # This does not pad
-                    T.RandomRotate(90, expand=True),
+                    T.RandomApply(T.RandomRotate(90, expand=True), 0.5),
                     T.Resize((args.input_size, args.input_size), preserve_aspect_ratio=True, symmetric_pad=True),
                 ]
                 if args.rotation

diff --git a/tests/common/test_transforms.py b/tests/common/test_transforms.py
@@ -27,6 +27,9 @@ def test_randomapply():
     transfo = T.RandomApply(lambda x: 1 - x)
     out = transfo(1)
     assert out == 0 or out == 1
+    transfo = T.RandomApply(lambda x, y: (1 - x, 2 * y))
+    out = transfo(1, np.array([2]))
+    assert out == (0, 4) or out == (1, 2) and isinstance(out[1], np.ndarray)
     assert repr(transfo).endswith(", p=0.5)")