1+ import cv2 as cv
2+ import numpy as np
3+ import tensorflow as tf
4+
5+ augmentations = [
6+ # 'additive_gaussian_noise',
7+ # 'additive_speckle_noise',
8+ 'random_brightness' ,
9+ 'random_contrast' ,
10+ # 'additive_shade',
11+ 'motion_blur'
12+ ]
13+
14+ def additive_gaussian_noise (image , stddev_range = [5 , 95 ]):
15+ stddev = tf .random_uniform ((), * stddev_range )
16+ noise = tf .random_normal (tf .shape (image ), stddev = stddev )
17+ noisy_image = tf .clip_by_value (image + noise , 0 , 255 )
18+ return noisy_image
19+
20+
21+ def additive_speckle_noise (image , prob_range = [0.0 , 0.005 ]):
22+ prob = tf .random_uniform ((), * prob_range )
23+ sample = tf .random_uniform (tf .shape (image ))
24+ noisy_image = tf .where (sample <= prob , tf .zeros_like (image ), image )
25+ noisy_image = tf .where (sample >= (1. - prob ), 255. * tf .ones_like (image ), noisy_image )
26+ return noisy_image
27+
28+
29+ def random_brightness (image , max_abs_change = 50 ):
30+ return tf .clip_by_value (tf .image .random_brightness (image , max_abs_change ), 0 , 255 )
31+
32+
33+ def random_contrast (image , strength_range = [0.5 , 1.5 ]):
34+ return tf .clip_by_value (tf .image .random_contrast (image , * strength_range ), 0 , 255 )
35+
36+
37+ def additive_shade (image , nb_ellipses = 20 , transparency_range = [- 0.5 , 0.8 ],
38+ kernel_size_range = [250 , 350 ]):
39+
40+ def _py_additive_shade (img ):
41+ min_dim = min (img .shape [:2 ]) / 4
42+ mask = np .zeros (img .shape [:2 ], np .uint8 )
43+ for i in range (nb_ellipses ):
44+ ax = int (max (np .random .rand () * min_dim , min_dim / 5 ))
45+ ay = int (max (np .random .rand () * min_dim , min_dim / 5 ))
46+ max_rad = max (ax , ay )
47+ x = np .random .randint (max_rad , img .shape [1 ] - max_rad ) # center
48+ y = np .random .randint (max_rad , img .shape [0 ] - max_rad )
49+ angle = np .random .rand () * 90
50+ cv .ellipse (mask , (x , y ), (ax , ay ), angle , 0 , 360 , 255 , - 1 )
51+
52+ transparency = np .random .uniform (* transparency_range )
53+ kernel_size = np .random .randint (* kernel_size_range )
54+ if (kernel_size % 2 ) == 0 : # kernel_size has to be odd
55+ kernel_size += 1
56+ mask = cv .GaussianBlur (mask .astype (np .float32 ), (kernel_size , kernel_size ), 0 )
57+ shaded = img * (1 - transparency * mask [..., np .newaxis ]/ 255. )
58+ return np .clip (shaded , 0 , 255 )
59+
60+ shaded = tf .py_func (_py_additive_shade , [image ], tf .float32 )
61+ res = tf .reshape (shaded , tf .shape (image ))
62+ return res
63+
64+
65+ def motion_blur (image , max_kernel_size = 10 ):
66+
67+ def _py_motion_blur (img ):
68+ # Either vertial, hozirontal or diagonal blur
69+ mode = np .random .choice (['h' , 'v' , 'diag_down' , 'diag_up' ])
70+ ksize = np .random .randint (0 , (max_kernel_size + 1 )/ 2 )* 2 + 1 # make sure is odd
71+ center = int ((ksize - 1 )/ 2 )
72+ kernel = np .zeros ((ksize , ksize ))
73+ if mode == 'h' :
74+ kernel [center , :] = 1.
75+ elif mode == 'v' :
76+ kernel [:, center ] = 1.
77+ elif mode == 'diag_down' :
78+ kernel = np .eye (ksize )
79+ elif mode == 'diag_up' :
80+ kernel = np .flip (np .eye (ksize ), 0 )
81+ var = ksize * ksize / 16.
82+ grid = np .repeat (np .arange (ksize )[:, np .newaxis ], ksize , axis = - 1 )
83+ gaussian = np .exp (- (np .square (grid - center )+ np .square (grid .T - center ))/ (2. * var ))
84+ kernel *= gaussian
85+ kernel /= np .sum (kernel )
86+ img = cv .filter2D (img , - 1 , kernel )
87+ return img
88+
89+ blurred = tf .numpy_function (_py_motion_blur , [image ], tf .float32 )
90+ return tf .reshape (blurred , tf .shape (image ))
91+
92+ def online_augmentation (image , random_order = True ):
93+ primitives = augmentations
94+ config = {}
95+ config ['random_brightness' ] = {'max_abs_change' : 50 }
96+ config ['random_contrast' ] = {'strength_range' : [0.3 , 1.5 ]}
97+ config ['additive_gaussian_noise' ] = {'stddev_range' : [0 , 10 ]}
98+ config ['additive_speckle_noise' ] = {'prob_range' : [0 , 0.0035 ]}
99+ config ['additive_shade' ] = {'transparency_range' : [- 0.5 , 0.5 ], 'kernel_size_range' : [100 , 150 ]}
100+ config ['motion_blur' ] = {'max_kernel_size' : 3 }
101+
102+ with tf .name_scope ('online_augmentation' ):
103+ prim_configs = [config .get (p , {}) for p in primitives ]
104+
105+ indices = tf .range (len (primitives ))
106+ if random_order :
107+ indices = tf .random .shuffle (indices )
108+
109+ def step (i , image ):
110+ fn_pairs = [(tf .equal (indices [i ], j ), lambda p = p , c = c : getattr (photaug , p )(image , ** c ))
111+ for j , (p , c ) in enumerate (zip (primitives , prim_configs ))]
112+ image = tf .case (fn_pairs )
113+ return i + 1 , image
114+
115+ _ , aug_image = tf .while_loop (lambda i , image : tf .less (i , len (primitives )),
116+ step , [0 , image ], parallel_iterations = 1 )
117+
118+ return aug_image
0 commit comments