SimCLR/utils.py
2020-02-18 15:06:14 -03:00

23 lines
574 B
Python

import cv2
import numpy as np
class GaussianBlur(object):
def __init__(self, min=0.1, max=2.0, kernel_size=9):
self.min = min
self.max = max
self.kernel_size = kernel_size
def __call__(self, sample):
sample = np.array(sample)
# blue the image with a 50% chance
prob = np.random.random_sample()
if prob < 0.5:
sigma = (self.max - self.min) * np.random.random_sample() + self.min
sample = cv2.GaussianBlur(sample, (self.kernel_size, self.kernel_size), sigma)
return sample