Source code for kitcar_ml.utils.data_generation.augmentation_utils

import cv2
import numpy as np


[docs]def gaussian_noise(img: np.ndarray, sigma: float, mean: float) -> np.ndarray: """Add gaussian noise to image.""" gauss = cv2.randn(np.empty(img.shape), mean, sigma).astype(np.float32) img = cv2.add(img, gauss) img = np.clip(img, 0, 255) return img
[docs]def add_random_intensity_and_contrast( img: np.ndarray, sigma_all: float, mean_all: float, gauss_blur_all: float ): """Add intensity/contrast to image.""" a = np.ones(img.shape, dtype=np.float32) sigma = sigma_all * np.random.randn() + 1 a = a * sigma img = cv2.multiply(img, a) b = np.ones(img.shape, dtype=np.float32) mean = mean_all * np.random.randn() b *= mean img = cv2.add(img, b) img = np.clip(img, 0, 255) img = cv2.GaussianBlur( img, (gauss_blur_all, gauss_blur_all), 0, ) return img
[docs]def adapt_contrast_and_mean(img: np.ndarray, dev_bckg_img: float): """Adapt contrast/mean of img to background.""" dev_aug = np.std(img) alpha = max(0.3, dev_bckg_img / dev_aug) a = np.ones(img.shape, np.float32) * alpha img = cv2.multiply( img, a ) # mean adaption: uncomment meanstuff above and replace img with img img = np.clip(img, 0, 255) return img
[docs]def supersample(img: np.ndarray, s: float) -> np.ndarray: """Supersample (resize) given image.""" return cv2.resize( img, None, fx=s, fy=s, interpolation=cv2.INTER_LINEAR, )
[docs]def revert_supersampling(img: np.ndarray, s: float) -> np.ndarray: """Revert previous supersampling (resize) of given image.""" return supersample(img, 1 / s)
[docs]def add_object_to_background( background_img: np.ndarray, aug_img: np.ndarray, img_coords: np.ndarray, noise_sigma: float, noise_mean: float, gauss_blur_sign: int, bckg_img_stddev: float, hor_motion_blur_size: int, ) -> np.ndarray: """Blend artificial image into background image. Args: background_img: Image without the current generation. aug_img: Pattern that should be inserted into background image. img_coords: Coordinates where to insert the aug_img. noise_mean: Mean of the noise. noise_sigma: Stddev of the noise. gauss_blur_sign: Gaussian blur kernel size applied to aug_img. bckg_img_stddev: Stddev of noise applied to the background img. hor_motion_blur_size: Kernel size of motion blur applied in hor direction. Return: Background image with aug_img added to it. """ mask = aug_img[:, :, 3] # Prepare aug img (traffic sign images) aug_img = cv2.cvtColor(aug_img, cv2.COLOR_BGR2GRAY) aug_img = adapt_contrast_and_mean(aug_img, bckg_img_stddev) # Motion blur kernel_h = np.zeros((hor_motion_blur_size, hor_motion_blur_size)) kernel_h[int((hor_motion_blur_size - 1) / 2), :] = np.ones(hor_motion_blur_size) kernel_h /= hor_motion_blur_size # Apply the horizontal kernel. aug_img = cv2.filter2D(aug_img, -1, kernel_h) aug_img = cv2.GaussianBlur(aug_img, (gauss_blur_sign, gauss_blur_sign), 0) aug_img = gaussian_noise(aug_img, noise_sigma, noise_mean) # Transform image to size it has within the final image aug_img_height, aug_img_width = aug_img.shape pts1 = np.float32( [[0, 0], [aug_img_width, 0], [0, aug_img_height], [aug_img_width, aug_img_height]] ) pts2 = np.float32( [[img_coords[:, 0]], [img_coords[:, 1]], [img_coords[:, 2]], [img_coords[:, 3]]] ) M = cv2.getPerspectiveTransform(pts1, pts2) aug_img = cv2.warpPerspective( aug_img, M, (background_img.shape[1], background_img.shape[0]), flags=cv2.INTER_LINEAR, ) mask = cv2.warpPerspective( mask, M, (background_img.shape[1], background_img.shape[0]), flags=cv2.INTER_LINEAR, ) _, mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY) mask = mask / 256 # Add aug_img to background image! generated_img = mask * aug_img + (1 - mask) * background_img return generated_img