Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_image_only_transform(image, mask):
height, width = image.shape[:2]
with mock.patch.object(ImageOnlyTransform, "apply") as mocked_apply:
with mock.patch.object(ImageOnlyTransform, "get_params", return_value={"interpolation": cv2.INTER_LINEAR}):
aug = ImageOnlyTransform(p=1)
data = aug(image=image, mask=mask)
mocked_apply.assert_called_once_with(image, interpolation=cv2.INTER_LINEAR, cols=width, rows=height)
assert np.array_equal(data["mask"], mask)
def apply(self, img, color_shift=0.05, intensity=1.0, random_state=None, **params):
return F.iso_noise(img, color_shift, intensity, np.random.RandomState(random_state))
def get_params(self):
return {
"color_shift": random.uniform(self.color_shift[0], self.color_shift[1]),
"intensity": random.uniform(self.intensity[0], self.intensity[1]),
"random_state": random.randint(0, 65536),
}
def get_transform_init_args_names(self):
return ("intensity", "color_shift")
class CLAHE(ImageOnlyTransform):
"""Apply Contrast Limited Adaptive Histogram Equalization to the input image.
Args:
clip_limit (float or (float, float)): upper threshold value for contrast limiting.
If clip_limit is a single float value, the range will be (1, clip_limit). Default: (1, 4).
tile_grid_size ((int, int)): size of grid for histogram equalization. Default: (8, 8).
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8
"""
def __init__(self, clip_limit=4.0, tile_grid_size=(8, 8), always_apply=False, p=0.5):
def __init__(
self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0
):
super(Normalize, self).__init__(always_apply, p)
self.mean = mean
self.std = std
self.max_pixel_value = max_pixel_value
def apply(self, image, **params):
return F.normalize(image, self.mean, self.std, self.max_pixel_value)
def get_transform_init_args_names(self):
return ("mean", "std", "max_pixel_value")
class Cutout(ImageOnlyTransform):
"""CoarseDropout of the square regions in the image.
Args:
num_holes (int): number of regions to zero out
max_h_size (int): maximum height of the hole
max_w_size (int): maximum width of the hole
fill_value (int, float, lisf of int, list of float): value for dropped pixels.
Targets:
image
Image types:
uint8, float32
Reference:
| https://arxiv.org/abs/1708.04552
def __init__(self, gamma_limit=(80, 120), eps=None, always_apply=False, p=0.5):
super(RandomGamma, self).__init__(always_apply, p)
self.gamma_limit = to_tuple(gamma_limit)
self.eps = eps
def apply(self, img, gamma=1, **params):
return F.gamma_transform(img, gamma=gamma)
def get_params(self):
return {"gamma": random.randint(self.gamma_limit[0], self.gamma_limit[1]) / 100.0}
def get_transform_init_args_names(self):
return ("gamma_limit", "eps")
class ToGray(ImageOnlyTransform):
"""Convert the input RGB image to grayscale. If the mean pixel value for the resulting image is greater
than 127, invert the resulting grayscale image.
Args:
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8, float32
"""
def apply(self, img, **params):
return F.to_gray(img)
def __init__(self, clip_limit=4.0, tile_grid_size=(8, 8), always_apply=False, p=0.5):
super(CLAHE, self).__init__(always_apply, p)
self.clip_limit = to_tuple(clip_limit, 1)
self.tile_grid_size = tuple(tile_grid_size)
def apply(self, img, clip_limit=2, **params):
return F.clahe(img, clip_limit, self.tile_grid_size)
def get_params(self):
return {"clip_limit": random.uniform(self.clip_limit[0], self.clip_limit[1])}
def get_transform_init_args_names(self):
return ("clip_limit", "tile_grid_size")
class ChannelDropout(ImageOnlyTransform):
"""Randomly Drop Channels in the input Image.
Args:
channel_drop_range (int, int): range from which we choose the number of channels to drop.
fill_value (int, float): pixel value for the dropped channel.
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8, uint16, unit32, float32
"""
def __init__(self, channel_drop_range=(1, 1), fill_value=0, always_apply=False, p=0.5):
super(ChannelDropout, self).__init__(always_apply, p)
p : float [0, 1], optional (default: 1.0)
Probability that the augmentation is performed to each image. Defaults
to ``1.0``.
"""
def __init__(self, idx, axis=1, always_apply=False, p=1.0):
super().__init__(always_apply, p)
self.idx = idx
self.axis = axis
def apply(self, im_arr, **params):
return np.delete(im_arr, self.idx, self.axis)
class SwapChannels(ImageOnlyTransform):
"""Swap channels in an input image.
Arguments
---------
first_idx : int
The first channel in the pair to swap.
second_idx : int
The second channel in the pair to swap.
axis : int, optional (default: 1)
The axis to drop the channel from. Defaults to ``1`` (torch channel
axis). Set to ``3`` for TF models where the channel is the last axis
of an image.
always_apply : bool, optional (default: False)
Apply this transformation to every image? Defaults to no (``False``).
p : float [0, 1], optional (default: 1.0)
Probability that the augmentation is performed to each image. Defaults
var = random.uniform(self.var_limit[0], self.var_limit[1])
sigma = var ** 0.5
random_state = np.random.RandomState(random.randint(0, 2 ** 32 - 1))
gauss = random_state.normal(self.mean, sigma, image.shape)
return {"gauss": gauss}
@property
def targets_as_params(self):
return ["image"]
def get_transform_init_args_names(self):
return ("var_limit",)
class ISONoise(ImageOnlyTransform):
"""
Apply camera sensor noise.
Args:
color_shift (float, float): variance range for color hue change.
Measured as a fraction of 360 degree Hue angle in HLS colorspace.
intensity ((float, float): Multiplicative factor that control strength
of color and luminace noise.
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8
"""
def __init__(self, quality_lower=99, quality_upper=100, always_apply=False, p=0.5):
super(JpegCompression, self).__init__(
quality_lower=quality_lower,
quality_upper=quality_upper,
compression_type=ImageCompression.ImageCompressionType.JPEG,
always_apply=always_apply,
p=p,
)
warnings.warn("This class has been deprecated. Please use ImageCompression", DeprecationWarning)
def get_transform_init_args(self):
return {"quality_lower": self.quality_lower, "quality_upper": self.quality_upper}
class RandomSnow(ImageOnlyTransform):
"""Bleach out some pixel values simulating snow.
From https://github.com/UjjwalSaxena/Automold--Road-Augmentation-Library
Args:
snow_point_lower (float): lower_bond of the amount of snow. Should be in [0, 1] range
snow_point_upper (float): upper_bond of the amount of snow. Should be in [0, 1] range
brightness_coeff (float): larger number will lead to a more snow on the image. Should be >= 0
Targets:
image
Image types:
uint8, float32
"""
def get_params_dependent_on_targets(self, params):
if not callable(self.mask):
return {"mask": self.mask}
return {"mask": self.mask(**params)}
@property
def targets_as_params(self):
return ["image"] + list(self.mask_params)
def get_transform_init_args_names(self):
return ("mode", "by_channels")
class RGBShift(ImageOnlyTransform):
"""Randomly shift values for each channel of the input RGB image.
Args:
r_shift_limit ((int, int) or int): range for changing values for the red channel. If r_shift_limit is a single
int, the range will be (-r_shift_limit, r_shift_limit). Default: (-20, 20).
g_shift_limit ((int, int) or int): range for changing values for the green channel. If g_shift_limit is a
single int, the range will be (-g_shift_limit, g_shift_limit). Default: (-20, 20).
b_shift_limit ((int, int) or int): range for changing values for the blue channel. If b_shift_limit is a single
int, the range will be (-b_shift_limit, b_shift_limit). Default: (-20, 20).
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8, float32
return ["image"]
def apply(self, img, channels_shuffled=[0, 1, 2], **params):
return F.channel_shuffle(img, channels_shuffled)
def get_params_dependent_on_targets(self, params):
img = params["image"]
ch_arr = list(range(img.shape[2]))
random.shuffle(ch_arr)
return {"channels_shuffled": ch_arr}
def get_transform_init_args_names(self):
return ()
class InvertImg(ImageOnlyTransform):
"""Invert the input image by subtracting pixel values from 255.
Args:
p (float): probability of applying the transform. Default: 0.5.
Targets:
image
Image types:
uint8
"""
def apply(self, img, **params):
return F.invert(img)
def get_transform_init_args_names(self):