Refactor core. Add typings. (#1219)

albumentations/augmentations/__init__.py

@@ -1,24 +1,19 @@
 # Common classes
-from .bbox_utils import *
 from .crops.functional import *
 from .crops.transforms import *
 
 # New transformations goes to individual files listed below
 from .domain_adaptation import *
-from .functional import *
-
-from .geometric.functional import *
-from .geometric.resize import *
-from .geometric.rotate import *
-from .geometric.transforms import *
-
 from .dropout.channel_dropout import *
 from .dropout.coarse_dropout import *
 from .dropout.cutout import *
 from .dropout.functional import *
 from .dropout.grid_dropout import *
 from .dropout.mask_dropout import *
-
-from .keypoints_utils import *
+from .functional import *
+from .geometric.functional import *
+from .geometric.resize import *
+from .geometric.rotate import *
+from .geometric.transforms import *
 from .transforms import *
 from .utils import *

albumentations/augmentations/crops/functional.py

@@ -1,18 +1,15 @@
-from typing import List, Optional, Sequence, Tuple, Union
+import typing
+from typing import Optional, Sequence, Tuple, Union
 
 import cv2
 import numpy as np
 
-from ..bbox_utils import denormalize_bbox, normalize_bbox
+from ...core.bbox_utils import denormalize_bbox, normalize_bbox
+from ...core.transforms_interface import BoxType, KeypointType
 from ..functional import _maybe_process_in_chunks, pad_with_params, preserve_channel_dim
 from ..geometric import functional as FGeometric
 
-BboxType = Union[List[int], List[float], Tuple[int, ...], Tuple[float, ...], np.ndarray]
-KeypointType = Union[List[int], List[float], Tuple[int, ...], Tuple[float, ...], np.ndarray]
-
 __all__ = [
-    "BboxType",
-    "KeypointType",
     "get_random_crop_coords",
     "random_crop",
     "crop_bbox_by_coords",
@@ -58,7 +55,7 @@ def random_crop(img: np.ndarray, crop_height: int, crop_width: int, h_start: flo
 
 
 def crop_bbox_by_coords(
-    bbox: BboxType, crop_coords: Tuple[int, int, int, int], crop_height: int, crop_width: int, rows: int, cols: int
+    bbox: BoxType, crop_coords: Tuple[int, int, int, int], crop_height: int, crop_width: int, rows: int, cols: int
 ):
     """Crop a bounding box using the provided coordinates of bottom-left and top-right corners in pixels and the
     required height and width of the crop.
@@ -75,15 +72,15 @@ def crop_bbox_by_coords(
         tuple: A cropped bounding box `(x_min, y_min, x_max, y_max)`.
 
     """
-    bbox = denormalize_bbox(bbox, rows, cols)
+    bbox = typing.cast(BoxType, denormalize_bbox(bbox, rows, cols))
     x_min, y_min, x_max, y_max = bbox[:4]
     x1, y1, _, _ = crop_coords
     cropped_bbox = x_min - x1, y_min - y1, x_max - x1, y_max - y1
     return normalize_bbox(cropped_bbox, crop_height, crop_width)
 
 
 def bbox_random_crop(
-    bbox: BboxType, crop_height: int, crop_width: int, h_start: float, w_start: float, rows: int, cols: int
+    bbox: BoxType, crop_height: int, crop_width: int, h_start: float, w_start: float, rows: int, cols: int
 ):
     crop_coords = get_random_crop_coords(rows, cols, crop_height, crop_width, h_start, w_start)
     return crop_bbox_by_coords(bbox, crop_coords, crop_height, crop_width, rows, cols)
@@ -150,7 +147,7 @@ def center_crop(img: np.ndarray, crop_height: int, crop_width: int):
     return img
 
 
-def bbox_center_crop(bbox: BboxType, crop_height: int, crop_width: int, rows: int, cols: int):
+def bbox_center_crop(bbox: BoxType, crop_height: int, crop_width: int, rows: int, cols: int):
     crop_coords = get_center_crop_coords(rows, cols, crop_height, crop_width)
     return crop_bbox_by_coords(bbox, crop_coords, crop_height, crop_width, rows, cols)
 
@@ -195,7 +192,7 @@ def crop(img: np.ndarray, x_min: int, y_min: int, x_max: int, y_max: int):
     return img[y_min:y_max, x_min:x_max]
 
 
-def bbox_crop(bbox: BboxType, x_min: int, y_min: int, x_max: int, y_max: int, rows: int, cols: int):
+def bbox_crop(bbox: BoxType, x_min: int, y_min: int, x_max: int, y_max: int, rows: int, cols: int):
     """Crop a bounding box.
 
     Args:
@@ -255,15 +252,15 @@ def crop_and_pad(
 
 
 def crop_and_pad_bbox(
-    bbox: Sequence[float],
+    bbox: BoxType,
     crop_params: Optional[Sequence[int]],
     pad_params: Optional[Sequence[int]],
     rows,
     cols,
     result_rows,
     result_cols,
     keep_size: bool,
-) -> Sequence[float]:
+) -> BoxType:
     x1, y1, x2, y2 = denormalize_bbox(bbox, rows, cols)
 
     if crop_params is not None:
@@ -273,21 +270,21 @@ def crop_and_pad_bbox(
         top, bottom, left, right = pad_params
         x1, y1, x2, y2 = x1 + left, y1 + top, x2 + left, y2 + top
 
-    bbox = normalize_bbox((x1, y1, x2, y2), result_rows, result_cols)
+    bbox = typing.cast(BoxType, normalize_bbox((x1, y1, x2, y2), result_rows, result_cols))
 
     return bbox
 
 
 def crop_and_pad_keypoint(
-    keypoint: Sequence[float],
+    keypoint: KeypointType,
     crop_params: Optional[Sequence[int]],
     pad_params: Optional[Sequence[int]],
     rows: int,
     cols: int,
     result_rows: int,
     result_cols: int,
     keep_size: bool,
-) -> Sequence[float]:
+) -> KeypointType:
     x, y, angle, scale = keypoint
 
     if crop_params is not None:

albumentations/augmentations/crops/transforms.py

@@ -5,8 +5,9 @@
 import cv2
 import numpy as np
 
-from ...core.transforms_interface import DualTransform, to_tuple
-from ..bbox_utils import union_of_bboxes
+from albumentations.core.bbox_utils import union_of_bboxes
+
+from ...core.transforms_interface import BoxType, DualTransform, KeypointType, to_tuple
 from ..geometric import functional as FGeometric
 from . import functional as F
 
@@ -686,28 +687,28 @@ def apply_to_mask(
 
     def apply_to_bbox(
         self,
-        bbox: Sequence[float],
+        bbox: BoxType,
         crop_params: Optional[Sequence[int]] = None,
         pad_params: Optional[Sequence[int]] = None,
         rows: int = 0,
         cols: int = 0,
         result_rows: int = 0,
         result_cols: int = 0,
         **params
-    ) -> Sequence[float]:
+    ) -> BoxType:
         return F.crop_and_pad_bbox(bbox, crop_params, pad_params, rows, cols, result_rows, result_cols, self.keep_size)
 
     def apply_to_keypoint(
         self,
-        keypoint: Sequence[float],
+        keypoint: KeypointType,
         crop_params: Optional[Sequence[int]] = None,
         pad_params: Optional[Sequence[int]] = None,
         rows: int = 0,
         cols: int = 0,
         result_rows: int = 0,
         result_cols: int = 0,
         **params
-    ) -> Sequence[float]:
+    ) -> KeypointType:
         return F.crop_and_pad_keypoint(
             keypoint, crop_params, pad_params, rows, cols, result_rows, result_cols, self.keep_size
         )

albumentations/augmentations/dropout/coarse_dropout.py

@@ -1,13 +1,12 @@
 import random
 import typing
-from typing import Optional, List, Tuple, Iterable, Union
+from typing import Iterable, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
-from ...core.transforms_interface import DualTransform
+from ...core.transforms_interface import DualTransform, KeypointType
 from .functional import cutout
 
-
 __all__ = ["CoarseDropout"]
 
 
@@ -161,19 +160,22 @@ def get_params_dependent_on_targets(self, params):
     def targets_as_params(self):
         return ["image"]
 
-    def _keypoint_in_hole(self, keypoint: Tuple[float, ...], hole: Tuple[int, int, int, int]) -> bool:
+    def _keypoint_in_hole(self, keypoint: KeypointType, hole: Tuple[int, int, int, int]) -> bool:
         x1, y1, x2, y2 = hole
         x, y = keypoint[:2]
         return x1 <= x < x2 and y1 <= y < y2
 
-    def apply_to_keypoints(self, keypoints: List[Tuple[float, ...]], holes: Iterable[Tuple] = (), **params):
+    def apply_to_keypoints(
+        self, keypoints: Sequence[KeypointType], holes: Iterable[Tuple[int, int, int, int]] = (), **params
+    ) -> List[KeypointType]:
+        result = []
         for hole in holes:
             remaining_keypoints = []
             for kp in keypoints:
-                if not self._keypoint_in_hole(kp, typing.cast(Tuple[int, int, int, int], hole)):
+                if not self._keypoint_in_hole(kp, hole):
                     remaining_keypoints.append(kp)
-            keypoints = remaining_keypoints
-        return keypoints
+            result = remaining_keypoints
+        return result
 
     def get_transform_init_args_names(self):
         return (

albumentations/augmentations/dropout/mask_dropout.py

@@ -1,12 +1,11 @@
 import random
-from typing import Union, Tuple, Any, Dict
+from typing import Any, Dict, Tuple, Union
 
 import cv2
 import numpy as np
 from skimage.measure import label
 
-from ...core.transforms_interface import DualTransform
-from ...core.transforms_interface import to_tuple
+from ...core.transforms_interface import DualTransform, to_tuple
 
 __all__ = ["MaskDropout"]
 
@@ -59,7 +58,7 @@ def get_params_dependent_on_targets(self, params) -> Dict[str, Any]:
         if num_labels == 0:
             dropout_mask = None
         else:
-            objects_to_drop = random.randint(self.max_objects[0], self.max_objects[1])
+            objects_to_drop = random.randint(int(self.max_objects[0]), int(self.max_objects[1]))
             objects_to_drop = min(num_labels, objects_to_drop)
 
             if objects_to_drop == num_labels:

albumentations/augmentations/functional.py

@@ -11,8 +11,7 @@
 import skimage
 
 from albumentations import random_utils
-
-from .keypoints_utils import angle_to_2pi_range
+from albumentations.core.keypoints_utils import angle_to_2pi_range
 
 __all__ = [
     "MAX_VALUES_BY_DTYPE",

albumentations/augmentations/geometric/functional.py

@@ -1,4 +1,5 @@
 import math
+import typing
 from typing import List, Optional, Sequence, Tuple, Union
 
 import cv2
@@ -7,7 +8,8 @@
 from scipy.ndimage.filters import gaussian_filter
 
 from ... import random_utils
-from ..bbox_utils import denormalize_bbox, normalize_bbox
+from ...core.bbox_utils import denormalize_bbox, normalize_bbox
+from ...core.transforms_interface import BoxType, KeypointType
 from ..functional import (
     _maybe_process_in_chunks,
     angle_2pi_range,
@@ -527,10 +529,10 @@ def warp_affine(
 
 @angle_2pi_range
 def keypoint_affine(
-    keypoint: Sequence[float],
+    keypoint: KeypointType,
     matrix: skimage.transform.ProjectiveTransform,
     scale: dict,
-) -> Sequence[float]:
+) -> KeypointType:
     if _is_identity_matrix(matrix):
         return keypoint
 
@@ -542,12 +544,12 @@ def keypoint_affine(
 
 
 def bbox_affine(
-    bbox: Sequence[float],
+    bbox: BoxType,
     matrix: skimage.transform.ProjectiveTransform,
     rows: int,
     cols: int,
     output_shape: Sequence[int],
-) -> Sequence[float]:
+) -> BoxType:
     if _is_identity_matrix(matrix):
         return bbox
 
@@ -568,7 +570,7 @@ def bbox_affine(
     y_min = np.min(points[:, 1])
     y_max = np.max(points[:, 1])
 
-    return normalize_bbox((x_min, y_min, x_max, y_max), output_shape[0], output_shape[1])
+    return typing.cast(BoxType, normalize_bbox((x_min, y_min, x_max, y_max), output_shape[0], output_shape[1]))
 
 
 @preserve_channel_dim
@@ -591,9 +593,7 @@ def safe_rotate(
     return warp_fn(img)
 
 
-def bbox_safe_rotate(
-    bbox: Tuple[float, float, float, float], matrix: np.ndarray, cols: int, rows: int
-) -> Tuple[float, float, float, float]:
+def bbox_safe_rotate(bbox: BoxType, matrix: np.ndarray, cols: int, rows: int) -> BoxType:
     x1, y1, x2, y2 = denormalize_bbox(bbox, rows, cols)
     points = np.array(
         [
@@ -620,18 +620,18 @@ def fix_point(pt1: float, pt2: float, max_val: float) -> Tuple[float, float]:
     x1, x2 = fix_point(x1, x2, cols)
     y1, y2 = fix_point(y1, y2, rows)
 
-    return normalize_bbox((x1, y1, x2, y2), rows, cols)
+    return typing.cast(BoxType, normalize_bbox((x1, y1, x2, y2), rows, cols))
 
 
 def keypoint_safe_rotate(
-    keypoint: Tuple[float, float, float, float],
+    keypoint: KeypointType,
     matrix: np.ndarray,
     angle: float,
     scale_x: float,
     scale_y: float,
     cols: int,
     rows: int,
-) -> Tuple[float, float, float, float]:
+) -> KeypointType:
     x, y, a, s = keypoint
     point = np.array([[x, y, 1]])
     x, y = (point @ matrix.T)[0]
@@ -816,4 +816,4 @@ def bbox_piecewise_affine(
     y1 = keypoints_arr[:, 1].min()
     x2 = keypoints_arr[:, 0].max()
     y2 = keypoints_arr[:, 1].max()
-    return normalize_bbox((x1, y1, x2, y2), h, w)
+    return typing.cast(BoxType, normalize_bbox((x1, y1, x2, y2), h, w))

albumentations/augmentations/geometric/resize.py

@@ -4,7 +4,7 @@
 import cv2
 import numpy as np
 
-from ...core.transforms_interface import DualTransform, to_tuple
+from ...core.transforms_interface import BoxType, DualTransform, KeypointType, to_tuple
 from . import functional as F
 
 __all__ = ["RandomScale", "LongestMaxSize", "SmallestMaxSize", "Resize"]
@@ -84,11 +84,11 @@ def apply(
     ) -> np.ndarray:
         return F.longest_max_size(img, max_size=max_size, interpolation=interpolation)
 
-    def apply_to_bbox(self, bbox: Sequence[float], **params) -> Sequence[float]:
+    def apply_to_bbox(self, bbox: BoxType, **params) -> BoxType:
         # Bounding box coordinates are scale invariant
         return bbox
 
-    def apply_to_keypoint(self, keypoint: Sequence[float], max_size: int = 1024, **params) -> Sequence[float]:
+    def apply_to_keypoint(self, keypoint: KeypointType, max_size: int = 1024, **params) -> KeypointType:
         height = params["rows"]
         width = params["cols"]
 
@@ -134,10 +134,10 @@ def apply(
     ) -> np.ndarray:
         return F.smallest_max_size(img, max_size=max_size, interpolation=interpolation)
 
-    def apply_to_bbox(self, bbox: Sequence[float], **params) -> Sequence[float]:
+    def apply_to_bbox(self, bbox: BoxType, **params) -> BoxType:
         return bbox
 
-    def apply_to_keypoint(self, keypoint: Sequence[float], max_size: int = 1024, **params) -> Sequence[float]:
+    def apply_to_keypoint(self, keypoint: KeypointType, max_size: int = 1024, **params) -> KeypointType:
         height = params["rows"]
         width = params["cols"]