Skip to content

Commit

Permalink
Move common functions to utils.py (#1260)
Browse files Browse the repository at this point in the history
* Move common functions into util.py

* Fix mypy errors
  • Loading branch information
Dipet committed Aug 24, 2022
1 parent a28dbb8 commit c3cb70a
Show file tree
Hide file tree
Showing 14 changed files with 453 additions and 369 deletions.
57 changes: 37 additions & 20 deletions albumentations/augmentations/crops/functional.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
import typing
from typing import Optional, Sequence, Tuple, Union
from typing import Optional, Sequence, Tuple

import cv2
import numpy as np

from albumentations.augmentations.utils import (
_maybe_process_in_chunks,
preserve_channel_dim,
)

from ...core.bbox_utils import denormalize_bbox, normalize_bbox
from ...core.transforms_interface import BoxType, KeypointType
from ..functional import _maybe_process_in_chunks, preserve_channel_dim
from ...core.transforms_interface import BoxInternalType, KeypointInternalType
from ..geometric import functional as FGeometric

__all__ = [
Expand Down Expand Up @@ -55,7 +58,12 @@ def random_crop(img: np.ndarray, crop_height: int, crop_width: int, h_start: flo


def crop_bbox_by_coords(
bbox: BoxType, crop_coords: Tuple[int, int, int, int], crop_height: int, crop_width: int, rows: int, cols: int
bbox: BoxInternalType,
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.
Expand All @@ -80,13 +88,15 @@ def crop_bbox_by_coords(


def bbox_random_crop(
bbox: BoxType, crop_height: int, crop_width: int, h_start: float, w_start: float, rows: int, cols: int
bbox: BoxInternalType, 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)


def crop_keypoint_by_coords(keypoint: KeypointType, crop_coords: Tuple[int, int, int, int]): # skipcq: PYL-W0613
def crop_keypoint_by_coords(
keypoint: KeypointInternalType, crop_coords: Tuple[int, int, int, int]
): # skipcq: PYL-W0613
"""Crop a keypoint using the provided coordinates of bottom-left and top-right corners in pixels and the
required height and width of the crop.
Expand All @@ -104,7 +114,13 @@ def crop_keypoint_by_coords(keypoint: KeypointType, crop_coords: Tuple[int, int,


def keypoint_random_crop(
keypoint: KeypointType, crop_height: int, crop_width: int, h_start: float, w_start: float, rows: int, cols: int
keypoint: KeypointInternalType,
crop_height: int,
crop_width: int,
h_start: float,
w_start: float,
rows: int,
cols: int,
):
"""Keypoint random crop.
Expand Down Expand Up @@ -147,12 +163,12 @@ def center_crop(img: np.ndarray, crop_height: int, crop_width: int):
return img


def bbox_center_crop(bbox: BoxType, crop_height: int, crop_width: int, rows: int, cols: int):
def bbox_center_crop(bbox: BoxInternalType, 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)


def keypoint_center_crop(keypoint: KeypointType, crop_height: int, crop_width: int, rows: int, cols: int):
def keypoint_center_crop(keypoint: KeypointInternalType, crop_height: int, crop_width: int, rows: int, cols: int):
"""Keypoint center crop.
Args:
Expand Down Expand Up @@ -192,7 +208,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: BoxType, x_min: int, y_min: int, x_max: int, y_max: int, rows: int, cols: int):
def bbox_crop(bbox: BoxInternalType, x_min: int, y_min: int, x_max: int, y_max: int, rows: int, cols: int):
"""Crop a bounding box.
Args:
Expand Down Expand Up @@ -230,9 +246,9 @@ def clamping_crop(img: np.ndarray, x_min: int, y_min: int, x_max: int, y_max: in
@preserve_channel_dim
def crop_and_pad(
img: np.ndarray,
crop_params: Sequence[int],
pad_params: Sequence[int],
pad_value: Union[int, float],
crop_params: Optional[Sequence[int]],
pad_params: Optional[Sequence[int]],
pad_value: Optional[float],
rows: int,
cols: int,
interpolation: int,
Expand All @@ -242,7 +258,9 @@ def crop_and_pad(
if crop_params is not None and any(i != 0 for i in crop_params):
img = crop(img, *crop_params)
if pad_params is not None and any(i != 0 for i in pad_params):
img = FGeometric.pad_with_params(img, *pad_params, border_mode=pad_mode, value=pad_value)
img = FGeometric.pad_with_params(
img, pad_params[0], pad_params[1], pad_params[2], pad_params[3], border_mode=pad_mode, value=pad_value
)

if keep_size:
resize_fn = _maybe_process_in_chunks(cv2.resize, dsize=(cols, rows), interpolation=interpolation)
Expand All @@ -252,15 +270,14 @@ def crop_and_pad(


def crop_and_pad_bbox(
bbox: BoxType,
bbox: BoxInternalType,
crop_params: Optional[Sequence[int]],
pad_params: Optional[Sequence[int]],
rows,
cols,
result_rows,
result_cols,
keep_size: bool,
) -> BoxType:
) -> BoxInternalType:
x1, y1, x2, y2 = denormalize_bbox(bbox, rows, cols)[:4]

if crop_params is not None:
Expand All @@ -274,15 +291,15 @@ def crop_and_pad_bbox(


def crop_and_pad_keypoint(
keypoint: KeypointType,
keypoint: KeypointInternalType,
crop_params: Optional[Sequence[int]],
pad_params: Optional[Sequence[int]],
rows: int,
cols: int,
result_rows: int,
result_cols: int,
keep_size: bool,
) -> KeypointType:
) -> KeypointInternalType:
x, y, angle, scale = keypoint[:4]

if crop_params is not None:
Expand Down
35 changes: 20 additions & 15 deletions albumentations/augmentations/crops/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,12 @@

from albumentations.core.bbox_utils import union_of_bboxes

from ...core.transforms_interface import BoxType, DualTransform, KeypointType, to_tuple
from ...core.transforms_interface import (
BoxInternalType,
DualTransform,
KeypointInternalType,
to_tuple,
)
from ..geometric import functional as FGeometric
from . import functional as F

Expand Down Expand Up @@ -450,7 +455,7 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, i

return {"x_min": x_min, "x_max": x_max, "y_min": y_min, "y_max": y_max}

def apply_to_bbox(self, bbox: BoxType, **params) -> BoxType:
def apply_to_bbox(self, bbox: BoxInternalType, **params) -> BoxInternalType:
return F.bbox_crop(bbox, **params)

def apply_to_keypoint(
Expand Down Expand Up @@ -678,11 +683,11 @@ def __init__(
def apply(
self,
img: np.ndarray,
crop_params: Sequence[int] = None,
pad_params: Sequence[int] = None,
pad_value: Union[int, float] = None,
rows: int = None,
cols: int = None,
crop_params: Sequence[int] = (),
pad_params: Sequence[int] = (),
pad_value: Union[int, float] = 0,
rows: int = 0,
cols: int = 0,
interpolation: int = cv2.INTER_LINEAR,
**params
) -> np.ndarray:
Expand All @@ -695,9 +700,9 @@ def apply_to_mask(
img: np.ndarray,
crop_params: Optional[Sequence[int]] = None,
pad_params: Optional[Sequence[int]] = None,
pad_value_mask: Union[int, float] = None,
rows: int = None,
cols: int = None,
pad_value_mask: float = None,
rows: int = 0,
cols: int = 0,
interpolation: int = cv2.INTER_NEAREST,
**params
) -> np.ndarray:
Expand All @@ -707,28 +712,28 @@ def apply_to_mask(

def apply_to_bbox(
self,
bbox: BoxType,
bbox: BoxInternalType,
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
) -> BoxType:
return F.crop_and_pad_bbox(bbox, crop_params, pad_params, rows, cols, result_rows, result_cols, self.keep_size)
) -> BoxInternalType:
return F.crop_and_pad_bbox(bbox, crop_params, pad_params, rows, cols, result_rows, result_cols)

def apply_to_keypoint(
self,
keypoint: KeypointType,
keypoint: KeypointInternalType,
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
) -> KeypointType:
) -> KeypointInternalType:
return F.crop_and_pad_keypoint(
keypoint, crop_params, pad_params, rows, cols, result_rows, result_cols, self.keep_size
)
Expand Down
7 changes: 4 additions & 3 deletions albumentations/augmentations/domain_adaptation.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,16 @@
from sklearn.decomposition import PCA
from sklearn.preprocessing import MinMaxScaler, StandardScaler

from ..core.transforms_interface import ImageOnlyTransform, to_tuple
from .functional import (
from albumentations.augmentations.utils import (
clipped,
get_opencv_dtype_from_numpy,
is_grayscale_image,
is_multispectral_image,
preserve_shape,
read_rgb_image,
)
from .utils import read_rgb_image

from ..core.transforms_interface import ImageOnlyTransform, to_tuple

__all__ = [
"HistogramMatching",
Expand Down
2 changes: 1 addition & 1 deletion albumentations/augmentations/dropout/functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import numpy as np

from ..functional import preserve_shape
from albumentations.augmentations.utils import preserve_shape

__all__ = ["cutout", "channel_dropout"]

Expand Down

0 comments on commit c3cb70a

Please sign in to comment.