62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.fft as fft
|
|
|
|
|
|
def Fourier_filter(x_in, threshold, scale):
|
|
x = x_in
|
|
B, C, H, W = x.shape
|
|
|
|
# Non-power of 2 images must be float32
|
|
if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
|
|
x = x.to(dtype=torch.float32)
|
|
|
|
# FFT
|
|
x_freq = fft.fftn(x, dim=(-2, -1))
|
|
x_freq = fft.fftshift(x_freq, dim=(-2, -1))
|
|
|
|
B, C, H, W = x_freq.shape
|
|
mask = torch.ones((B, C, H, W), device=x.device)
|
|
|
|
crow, ccol = H // 2, W // 2
|
|
mask[..., crow - threshold: crow + threshold,
|
|
ccol - threshold: ccol + threshold] = scale
|
|
x_freq = x_freq * mask
|
|
|
|
# IFFT
|
|
x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
|
|
x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
|
|
|
|
return x_filtered.to(dtype=x_in.dtype)
|
|
|
|
|
|
class ScaleU(nn.Module):
|
|
def __init__(self, enable_se_scaleu=True, b_size=1280, s_size=1):
|
|
super(ScaleU, self).__init__()
|
|
self.scaleu_b = nn.Parameter(torch.zeros(b_size))
|
|
self.scaleu_s = nn.Parameter(torch.zeros(s_size))
|
|
self.enable_se_scaleu = enable_se_scaleu
|
|
|
|
def forward(self, h, hs_, transformer_options={}):
|
|
h = h.to(torch.float32)
|
|
hs_ = hs_.to(torch.float32)
|
|
b = torch.tanh(self.scaleu_b) + 1
|
|
s = torch.tanh(self.scaleu_s) + 1
|
|
if self.enable_se_scaleu:
|
|
hidden_mean = h.mean(1).unsqueeze(1) # B,1,H,W
|
|
B = hidden_mean.shape[0]
|
|
hidden_max, _ = torch.max(hidden_mean.view(
|
|
B, -1), dim=-1, keepdim=True) # B,1
|
|
hidden_min, _ = torch.min(hidden_mean.view(
|
|
B, -1), dim=-1, keepdim=True) # B,1
|
|
# duplicate the hidden_mean dimension 1 to C
|
|
hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / \
|
|
(hidden_max - hidden_min).unsqueeze(2).unsqueeze(3) # B,1,H,W
|
|
b = torch.einsum('c,bchw->bchw', b-1, hidden_mean) + 1.0 # B,C,H,W
|
|
h = torch.einsum('bchw,bchw->bchw', h, b)
|
|
else:
|
|
h = torch.einsum('bchw,c->bchw', h, b)
|
|
|
|
hs_ = Fourier_filter(hs_, threshold=1, scale=s)
|
|
return h.to(torch.float16), hs_.to(torch.float16)
|