1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
import numpy as np
import torch
from torchvision.transforms import transforms
def color_distortion(s=1.0):
# s is the strength of color distortion.
color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)
rnd_gray = transforms.RandomGrayscale(p=0.2)
color_distort = transforms.Compose([
rnd_color_jitter,
rnd_gray
])
return color_distort
class Clip(object):
def __call__(self, x):
return torch.clamp(x, 0, 1)
class RandomGaussianBlur(object):
"""
PyTorch version of
https://github.com/google-research/simclr/blob/244e7128004c5fd3c7805cf3135c79baa6c3bb96/data_util.py#L311
"""
def gaussian_blur(self, image, sigma):
image = image.reshape(1, 3, 224, 224)
radius = np.int(self.kernel_size / 2)
kernel_size = radius * 2 + 1
x = np.arange(-radius, radius + 1)
blur_filter = np.exp(
-np.power(x, 2.0) / (2.0 * np.power(np.float(sigma), 2.0)))
blur_filter /= np.sum(blur_filter)
conv1 = torch.nn.Conv2d(3, 3, kernel_size=(kernel_size, 1), groups=3,
padding=[kernel_size // 2, 0], bias=False)
conv1.weight = torch.nn.Parameter(torch.Tensor(np.tile(
blur_filter.reshape(kernel_size, 1, 1, 1), 3
).transpose([3, 2, 0, 1])))
conv2 = torch.nn.Conv2d(3, 3, kernel_size=(1, kernel_size), groups=3,
padding=[0, kernel_size // 2], bias=False)
conv2.weight = torch.nn.Parameter(torch.Tensor(np.tile(
blur_filter.reshape(kernel_size, 1, 1, 1), 3
).transpose([3, 2, 1, 0])))
res = conv2(conv1(image))
assert res.shape == image.shape
return res[0]
def __init__(self, kernel_size, sigma_range=(0.1, 2), p=0.5):
self.kernel_size = kernel_size
self.sigma_range = sigma_range
self.p = p
def __call__(self, img):
with torch.no_grad():
assert isinstance(img, torch.Tensor)
if np.random.uniform() < self.p:
return self.gaussian_blur(
img, sigma=np.random.uniform(*self.sigma_range)
)
return img
|
