--- a +++ b/Segmentation/utils.py @@ -0,0 +1,79 @@ +import random +import cv2 +import numpy as np + +def gaussian_noise(img, mean=0, sigma=0.03): + img = img.copy() + noise = np.random.normal(mean, sigma, img.shape) + mask_overflow_upper = img+noise >= 1.0 + mask_overflow_lower = img+noise < 0 + noise[mask_overflow_upper] = 1.0 + noise[mask_overflow_lower] = 0 + img += noise + return img + +def brightness(img, low, high): + value = random.uniform(low, high) + hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) + hsv = np.array(hsv, dtype = np.float64) + hsv[:,:,1] = hsv[:,:,1]*value + hsv[:,:,1][hsv[:,:,1]>255] = 255 + hsv[:,:,2] = hsv[:,:,2]*value + hsv[:,:,2][hsv[:,:,2]>255] = 255 + hsv = np.array(hsv, dtype = np.uint8) + img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + return img + +def fill(img, h, w): + img = img.astype('float32') + img = cv2.resize(img, (h, w), cv2.INTER_CUBIC) + return img +def zoom(img, mask, value): + if value > 1 or value < 0: + print('Value for zoom should be less than 1 and greater than 0') + return img, mask + value = random.uniform(value, 1) + h, w = img.shape[:2] + h_taken = int(value*h) + w_taken = int(value*w) + h_start = random.randint(0, h-h_taken) + w_start = random.randint(0, w-w_taken) + img = img[h_start:h_start+h_taken, w_start:w_start+w_taken, :] + img = fill(img, h, w) + mask = mask[h_start:h_start+h_taken, w_start:w_start+w_taken, :] + mask = fill(mask, h, w) + return img, mask + +def vertical_shift(img, mask, ratio=0.0): + if ratio > 1 or ratio < 0: + print('Value should be less than 1 and greater than 0') + return img, mask + ratio = random.uniform(-ratio, ratio) + h, w = img.shape[:2] + to_shift = h*ratio + if ratio > 0: + img = img[:, :int(w-to_shift), :] + mask = mask[:, :int(w-to_shift), :] + if ratio < 0: + img = img[:, int(-1*to_shift):, :] + mask = mask[:, int(-1*to_shift):, :] + img = fill(img, h, w) + mask = fill(mask, h, w) + return img, mask + +def horizontal_shift(img, mask, ratio=0.0): + if ratio > 1 or ratio < 0: + print('Value should be less than 1 and greater than 0') + return img, mask + ratio = random.uniform(-ratio, ratio) + h, w = img.shape[:2] + to_shift = w*ratio + if ratio > 0: + img = img[:, :int(w-to_shift), :] + mask = mask[:, :int(w-to_shift), :] + if ratio < 0: + img = img[:, int(-1*to_shift):, :] + mask = mask[:, int(-1*to_shift):, :] + img = fill(img, h, w) + mask = fill(mask, h, w) + return img,mask \ No newline at end of file