From a8bd7066ded6df95ce58a1c2d0d48887db2daabd Mon Sep 17 00:00:00 2001
From: Kisuk Lee <kisuk@zetta.ai>
Date: Tue, 23 Jan 2024 22:13:55 +0900
Subject: [PATCH] feat(augment): add `recompute` and `border`

---
 .../data/augment/pinky_basil/aug_mip1_v3.py   | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/deepem/data/augment/pinky_basil/aug_mip1_v3.py b/deepem/data/augment/pinky_basil/aug_mip1_v3.py
index f6105d0..21b6afe 100644
--- a/deepem/data/augment/pinky_basil/aug_mip1_v3.py
+++ b/deepem/data/augment/pinky_basil/aug_mip1_v3.py
@@ -1,8 +1,17 @@
 from augmentor import *
 
 
-def get_augmentation(is_train, box=None, missing=7, blur=7, lost=True,
-                     random=False, **kwargs):
+def get_augmentation(
+    is_train,
+    box=None,
+    missing=7,
+    blur=7,
+    lost=True,
+    random=False,
+    recompute=False,
+    border=False,
+    **kwargs
+):
     augs = list()
 
     # Box
@@ -72,7 +81,15 @@ def get_augmentation(is_train, box=None, missing=7, blur=7, lost=True,
     if is_train:
         augs.append(Warp(skip=0.3, do_twist=False, rot_max=45.0, scale_max=1.1))
 
+    # Recompute connected components
+    if recompute:
+        augs.append(Label())
+
     # Flip & rotate
     augs.append(FlipRotate())
 
+    # Create border
+    if border:
+        augs.append(Border())
+
     return Compose(augs)