[orientation] augment angle while training (#1567)

mindee · Apr 26, 2024 · 5568612 · 5568612
1 parent 124960d
commit 5568612
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/references/classification/train_pytorch_orientation.py b/references/classification/train_pytorch_orientation.py
@@ -41,7 +41,10 @@
 def rnd_rotate(img: torch.Tensor, target):
     angle = int(np.random.choice(CLASSES))
     idx = CLASSES.index(angle)
-    rotated_img = F.rotate(img, angle=-angle, fill=0, expand=False)[:3]
+    # augment the angle randomly with a probability of 0.5
+    if np.random.rand() < 0.5:
+        angle += float(np.random.choice(np.arange(-25, 25, 5)))
+    rotated_img = F.rotate(img, angle=-angle, fill=0, expand=angle not in CLASSES)[:3]
     return rotated_img, idx
 
 

diff --git a/references/classification/train_tensorflow_orientation.py b/references/classification/train_tensorflow_orientation.py
@@ -36,8 +36,11 @@
 def rnd_rotate(img: tf.Tensor, target):
     angle = int(np.random.choice(CLASSES))
     idx = CLASSES.index(angle)
+    # augment the angle randomly with a probability of 0.5
+    if np.random.rand() < 0.5:
+        angle += float(np.random.choice(np.arange(-25, 25, 5)))
     # clockwise rotation
-    rotated_img = rotated_img_tensor(img, -angle, expand=False)
+    rotated_img = rotated_img_tensor(img, -angle, expand=angle not in CLASSES)
     return rotated_img, idx