|
|
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,96)): |
|
|
super().__init__() |
|
|
self.upsample = lambda x: F.interpolate(x, [224,224]) |
|
|
self.downsample = lambda x: F.interpolate(x, [resolution[0],resolution[1]]) |
|
|
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 |
