File size: 3,888 Bytes
caa56d6 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import cv2
import numpy as np
import math
import numbers
import torch
from torch import nn
from torch.nn import functional as F
class GaussianSmoothing(nn.Module):
"""
Apply gaussian smoothing on a
1d, 2d or 3d tensor. Filtering is performed seperately for each channel
in the input using a depthwise convolution.
Arguments:
channels (int, sequence): Number of channels of the input tensors. Output will
have this number of channels as well.
kernel_size (int, sequence): Size of the gaussian kernel.
sigma (float, sequence): Standard deviation of the gaussian kernel.
dim (int, optional): The number of dimensions of the data.
Default value is 2 (spatial).
"""
def __init__(self, channels, kernel_size, sigma=0.1, dim=2):
super(GaussianSmoothing, self).__init__()
self.kernel_size = kernel_size
if isinstance(kernel_size, numbers.Number):
kernel_size = [kernel_size] * dim
if isinstance(sigma, numbers.Number):
sigma = [sigma] * dim
# The gaussian kernel is the product of the
# gaussian function of each dimension.
kernel = 1
meshgrids = torch.meshgrid(
[
torch.arange(size, dtype=torch.float32)
for size in kernel_size
]
)
for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
mean = (size - 1) / 2
kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
torch.exp(-((mgrid - mean) / std) ** 2 / 2)
# Make sure sum of values in gaussian kernel equals 1.
kernel = kernel / torch.sum(kernel)
# Reshape to depthwise convolutional weight
kernel = kernel.view(1, 1, *kernel.size())
kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
self.register_buffer('weight', kernel)
self.groups = channels
if dim == 1:
self.conv = F.conv1d
elif dim == 2:
self.conv = F.conv2d
elif dim == 3:
self.conv = F.conv3d
else:
raise RuntimeError(
'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(
dim)
)
def forward(self, input):
"""
Apply gaussian filter to input.
Arguments:
input (torch.Tensor): Input to apply gaussian filter on.
Returns:
filtered (torch.Tensor): Filtered output.
"""
if self.training:
return self.conv(input, weight=self.weight, groups=self.groups, padding=self.kernel_size//2)
else:
return input
class GaussianNoise(nn.Module):
def __init__(self, mean=0, std=0.1, clip=1):
super(GaussianNoise, self).__init__()
self.mean = mean
self.std = std
self.clip = clip
def forward(self, x):
if self.training:
noise = x.data.new(x.size()).normal_(self.mean, self.std)
return torch.clamp(x + noise, -self.clip, self.clip)
else:
return x
if __name__ == "__main__":
im = cv2.imread('E:\SRM\component\FF-F2F_0.png')
im_ten = im/255*2-1
im_ten = torch.from_numpy(im_ten).unsqueeze(0).permute(0, 3, 1, 2).float()
blur = GaussianSmoothing(channels=3, kernel_size=7, sigma=0.8)
noise = GaussianNoise()
noise_im = torch.clamp(noise(im_ten), -1, 1)
blur_im = blur(im_ten)
print(blur_im.size())
def t2im(t):
t = (t+1)/2*255
im = t.squeeze().cpu().numpy().transpose(1, 2, 0).astype(np.uint8)
return im
cv2.imshow('ori', im)
cv2.imshow('blur', t2im(blur_im))
cv2.imshow('noise', t2im(noise_im))
cv2.waitKey()
|