from ..registry import FUNCTIONAL_REGISTRY
from .operations import *
[docs]@FUNCTIONAL_REGISTRY.register()
class FlowAugmentor:
"""
Class for appyling a series of augmentations to a pair of images and a flow field
Parameters
----------
crop_size : int
Size of the crop to be applied to the images.
eraser_aug_params : dict
Parameters for the eraser augmentation.
noise_aug_params : dict
Parameters for the noise augmentation.
flip_aug_params : dict
Parameters for the flip augmentation.
color_aug_params : dict
Parameters for the color augmentation.
spatial_aug_params : dict
Parameters for the spatial augmentation.
advanced_spatial_aug_params : dict
Parameters for the spatial augmentation.
If both spatial and advanced_spatial augmentations are enabled, the advanced spatial augmentations are used.
"""
def __init__(
self,
crop_size,
eraser_aug_params={"enabled": False, "aug_prob": 0.5, "bounds": [50, 100]},
noise_aug_params={"enabled": False, "aug_prob": 0.5, "noise_std_range": 0.06},
flip_aug_params={"enabled": False, "h_flip_prob": 0.5, "v_flip_prob": 0.1},
color_aug_params={
"enabled": False,
"asymmetric_color_aug_prob": 0.2,
"brightness": 0.4,
"contrast": 0.4,
"saturation": 0.4,
"hue": 0.15915494309189535,
},
spatial_aug_params={
"enabled": False,
"aug_prob": 0.8,
"stretch_prob": 0.8,
"min_scale": -0.1,
"max_scale": 1.0,
"max_stretch": 0.2,
},
advanced_spatial_aug_params={
"enabled": False,
"scale1": 0.3,
"scale2": 0.1,
"rotate": 0.4,
"translate": 0.4,
"stretch": 0.3,
"enable_out_of_boundary_crop": False,
},
):
self.crop_size = crop_size
self.color_aug_params = color_aug_params
self.eraser_aug_params = eraser_aug_params
self.spatial_aug_params = spatial_aug_params
self.noise_aug_params = noise_aug_params
self.flip_aug_params = flip_aug_params
self.advanced_spatial_aug_params = advanced_spatial_aug_params
self.advanced_spatial_aug_params["h_flip_prob"] = (
flip_aug_params["h_flip_prob"] if "h_flip_prob" in flip_aug_params else 0.0
)
self.advanced_spatial_transform = AdvancedSpatialTransform(
crop=self.crop_size, **self.advanced_spatial_aug_params
)
if self.advanced_spatial_aug_params["enabled"]:
# Disable spatial transform and horizontal flip if advanced spatial transforms are used
self.spatial_aug_params["enabled"] = False
self.flip_aug_params["h_flip_prob"] = 0.0
def __call__(self, img1, img2, flow, valid=None):
"""
Applies the augmentations to the pair of images and the flow field.
Parameters
----------
img1 : numpy.ndarray
First image
img2 : numpy.ndarray
Second image
flow : numpy.ndarray
Flow field
valid : default: None
None object
Returns
-------
img1 : numpy.ndarray
First image
img2 : numpy.ndarray
Second image
flow : numpy.ndarray
Flow field
valid : None
None object
"""
img1, img2 = color_transform(img1, img2, **self.color_aug_params)
img1, img2, flow = self.advanced_spatial_transform(img1, img2, flow)
img1, img2, flow = spatial_transform(
img1, img2, flow, self.crop_size, **self.spatial_aug_params
)
img1, img2, flow, _ = flip_transform(img1, img2, flow, **self.flip_aug_params)
img1, img2 = noise_transform(img1, img2, **self.noise_aug_params)
img1, img2 = eraser_transform(img1, img2, **self.eraser_aug_params)
img1 = np.ascontiguousarray(img1)
img2 = np.ascontiguousarray(img2)
flow = np.ascontiguousarray(flow)
return img1, img2, flow, None
[docs]@FUNCTIONAL_REGISTRY.register()
class SparseFlowAugmentor(FlowAugmentor):
"""
Class for appyling a series of augmentations to a pair of images, flow field and valid flow field.
Parameters
----------
crop_size : int
Size of the crop to be applied to the images.
color_aug_params : dict
Parameters for the color augmentation.
eraser_aug_params : dict
Parameters for the eraser augmentation.
spatial_aug_params : dict
Parameters for the spatial augmentation.
flip_aug_params : dict
Parameters for the flip augmentation.
"""
def __call__(self, img1, img2, flow, valid):
"""
Applies the augmentations to the pair of images and the flow field.
Parameters
----------
img1 : numpy.ndarray
First image
img2 : numpy.ndarray
Second image
flow : numpy.ndarray
Flow field
valid : numpy.ndarray
Valid Flow field
Returns
-------
img1 : numpy.ndarray
First image
img2 : numpy.ndarray
Second image
flow : numpy.ndarray
Flow field
valid : numpy.ndarray
Valid Flow field
"""
img1, img2 = color_transform(img1, img2, **self.color_aug_params)
img1, img2 = eraser_transform(img1, img2, **self.eraser_aug_params)
img1, img2, flow, valid = sparse_spatial_transform(
img1, img2, flow, valid, self.crop_size, **self.spatial_aug_params
)
img1, img2, flow, valid = flip_transform(
img1, img2, flow, valid, **self.flip_aug_params
)
img1 = np.ascontiguousarray(img1)
img2 = np.ascontiguousarray(img2)
flow = np.ascontiguousarray(flow)
valid = np.ascontiguousarray(valid)
return img1, img2, flow, valid