【数字图像处理】频域低通滤波

代码

思想:将图像进行FFT变换后，把DC分量移到图像中心，然后截取中间HxW的小矩阵（即保留低频分量，舍去高频分量）

import torch
import torch.fft
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch.nn as nn

class SpectralPooling(nn.Module):
    def __init__(self,h,w):
        super().__init__()
        self.h=h
        self.w=w

    def forward(self,x):
        batch,channel,height,width=x.size()
        pad_h=0
        pad_w=0
        left=(width-self.w)//2
        top=(height-self.h)//2
        if height%2==0:
            pad_h=1
            top=(height+1-self.h)//2
        if width%2==0:
            pad_w=1
            left=(width+1-self.w)//2
        x_padded=F.pad(input=x,pad=(0,pad_w,0,pad_h))
        x_ffted=torch.fft.fftshift(torch.fft.fftn(x_padded,dim=(-2,-1)))
        x_croped=x_ffted[:,:,top:top+self.h,left:left+self.w]
        y=torch.fft.ifftn(torch.fft.ifftshift(x_croped),dim=(-2,-1))
        return y

def spectralpool(x,h,w):
    batch,channel,height,width=x.size()
    pad_h=0
    pad_w=0
    left = (width - w) // 2
    top = (height - h) // 2
    if height%2==0:
        pad_h=1
        top=(height+1-h)//2
    if width%2==0:
        pad_w=1
        left=(width+1-w)//2
    x_padded=F.pad(input=x,pad=(0,pad_w,0,pad_h))
    x_ffted=torch.fft.fftshift(torch.fft.fftn(x_padded,dim=(-2,-1)))

    x_croped=x_ffted[:,:,top:top+h,left:left+w]
    y=torch.fft.ifftn(torch.fft.ifftshift(x_croped),dim=(-2,-1))
    return y

#检验共轭对称性是否保持
x=torch.randn((4,3,18,24))
y=spectralpool(x,7,11)
print(y.size())
print(torch.max(torch.abs(y.imag)))
#图像滤波实验
h=153
w=147
im = Image.open('img.png').convert('RGB')
im.show()
im = np.array(im)
x=torch.from_numpy(im).permute(2,0,1).contiguous().view(1,3,758,850)
spool=SpectralPooling(h,w)
y=spool.forward(x)
print(torch.max(torch.abs(y.imag/y.real)))
z=y.view(3,h,w).permute(1,2,0).real.numpy()
min=np.min(z)
max=np.max(z)
z=(z-min)/(max-min)
z=np.round(z*255)
pic = Image.fromarray(z.astype("uint8"))
plt.imshow(pic)
plt.show()

实验结果

图像滤波效果

滤波前:

滤波后: