Added spatial bbox helper

augly/image/utils/bboxes.py

@@ -2,12 +2,23 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
 
 import math
-from typing import List, Optional, Tuple
+from PIL import Image
+from typing import Callable, List, Optional, Tuple
 
 import augly.image.utils as imutils
 import numpy as np
 
 
+def create_test_image(w: int, h: int, bbox: Tuple) -> Image.Image:
+    """
+    Create dummy test image to help spatial_bbox_helper.
+    """
+    image = Image.new("RGB", (w, h))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle([bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h], fill="white")
+    return image
+
+
 def crop_bboxes_helper(
     bbox: Tuple, x1: float, y1: float, x2: float, y2: float, **kwargs
 ) -> Tuple:
@@ -444,6 +455,29 @@ def get_enclosing_bbox(
     )
 
 
+def spatial_bbox_helper(
+    bbox: Tuple, src_w: int, src_h: int, aug_function: Callable, **kwargs
+) -> Tuple:
+    """
+    Computes the bbox that encloses the transformed bbox in the image transformed by
+    `aug_function`. This helper can be used to compute the transformed bbox for any
+    augmentation which doesn't affect the color of the source image (e.g. any spatial
+    augmentation).
+    """
+    dummy_image = create_test_image(w=src_w, h=src_h, bbox=bbox)
+
+    aug_image = aug_function(dummy_image, **kwargs)
+    aug_w, aug_h = aug_image.size
+    array_image = np.array(aug_image)
+
+    white_y, white_x, _ = np.where(array_image > 0)
+    min_x, max_x = np.min(white_x), np.max(white_x)
+    min_y, max_y = np.min(white_y), np.max(white_y)
+
+    new_bbox = (min_x / aug_w, min_y / aug_h, max_x / aug_w, max_y / aug_h)
+    return new_bbox
+
+
 def vflip_bboxes_helper(bbox: Tuple, **kwargs) -> Tuple:
     """
     Analogous to hflip, when the src image is vertically flipped, the bounding box also