|
|
import torch.nn as nn
|
|
|
|
|
|
from torch.nn import functional as F
|
|
|
from timm import create_model
|
|
|
|
|
|
|
|
|
__all__ = ['NoiseTransformer']
|
|
|
|
|
|
class NoiseTransformer(nn.Module):
|
|
|
def __init__(self, resolution=128):
|
|
|
super().__init__()
|
|
|
self.upsample = lambda x: F.interpolate(x, [224,224])
|
|
|
self.downsample = lambda x: F.interpolate(x, [resolution,resolution])
|
|
|
self.upconv = nn.Conv2d(7,4,(1,1),(1,1),(0,0))
|
|
|
self.downconv = nn.Conv2d(4,3,(1,1),(1,1),(0,0))
|
|
|
|
|
|
self.swin = create_model("swin_tiny_patch4_window7_224",pretrained=True)
|
|
|
|
|
|
|
|
|
def forward(self, x, residual=False):
|
|
|
if residual:
|
|
|
x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x))))) + x
|
|
|
else:
|
|
|
x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x)))))
|
|
|
|
|
|
return x |
