[图像处理]基于 PyTorch 的高斯核卷积

Posted 2021-03-15 xxxxxxxxx

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()