Shortcuts

Source code for catalyst.data.cv.mixins.blur

from typing import List
import random

import numpy as np

import albumentations as albu


[docs]class BlurMixin: """Calculates blur factor for augmented image."""
[docs] def __init__( self, input_key: str = "image", output_key: str = "blur_factor", blur_min: int = 3, blur_max: int = 9, blur: List[str] = None, ): """ Args: input_key: input key to use from annotation dict output_key: output key to use to store the result """ self.input_key = input_key self.output_key = output_key self.blur_min = blur_min self.blur_max = blur_max blur = blur or ["Blur"] self.blur = [albu.__dict__[x]() for x in blur] self.num_blur = len(self.blur) self.num_blur_classes = blur_max - blur_min + 1 + 1 self.blur_probability = ( self.num_blur_classes - 1 ) / self.num_blur_classes
def __call__(self, dictionary): """@TODO: Docs. Contribution is welcome.""" image = dictionary[self.input_key] blur_factor = 0 if random.random() < self.blur_probability: blur_fn = np.random.choice(self.blur) blur_factor = int( np.random.randint(self.blur_min, self.blur_max) - self.blur_min + 1 ) image = blur_fn.apply(image=image, ksize=blur_factor) dictionary[self.input_key] = image dictionary[self.output_key] = blur_factor return dictionary
__all__ = ["BlurMixin"]