import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import cv2
import matplotlib.pyplot as plt
from PIL import Image
class GaussianBlurConv(nn.Module):
def __init__(self, channels=3):
super(GaussianBlurConv, self).__init__()
self.channels = channels
# print("channels: ", channels.shape)
kernel = [[0.00078633, 0.00655965, 0.01330373, 0.00655965, 0.00078633],
[0.00655965, 0.05472157, 0.11098164, 0.05472157, 0.00655965],
[0.01330373, 0.11098164, 0.22508352, 0.11098164, 0.01330373],
[0.00655965, 0.05472157, 0.11098164, 0.05472157, 0.00655965],
[0.00078633, 0.00655965, 0.01330373, 0.00655965, 0.00078633]]
kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0) # (H, W) -> (1, 1, H, W)
kernel = kernel.expand((int(channels), 1, 5, 5))
self.weight = nn.Parameter(data=kernel, requires_grad=False)
def __call__(self, x):
x = F.conv2d(x, self.weight, padding=2, groups=self.channels)
return x
path = r"/home/curry/Pictures/lenna/lenna.jpg"
img = Image.open(path)
img = np.array(img)
img = torch.from_numpy(img).unsqueeze(dim=0).permute(0, 3, 1, 2).contiguous().float() # (1, H, W, C) -> (1, C, H, W)
gaussian = GaussianBlurConv(channels=img.shape[1])
result = gaussian(img).squeeze().permute(1, 2, 0).contiguous().numpy().astype(np.int32) # 做高斯处理
plt.imshow(result)
plt.show()
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:[图像处理]基于 PyTorch 的高斯核卷积 - Python技术站