def __init__(self):
super().__init__()
self.model = torch.nn.Sequential(*[
Conv2d(in_channels=1,
out_channels=8,
kernel_size=3,
stride=1,
padding=1),
Conv2d(in_channels=8,
out_channels=16,
kernel_size=3,
stride=1,
padding=1),
Conv2d(in_channels=16,
out_channels=32,
kernel_size=3,
stride=1,
padding=1),
PixelShuffle(upscale_factor=2),
Conv2d(in_channels=8,
out_channels=1,
kernel_size=3,
stride=1,
padding=1),
])
def __init__(self, in_channels, up_scale):
super(UpSampleBlock, self).__init__()
self.conv = Conv2d(in_channels,
in_channels * up_scale**2,
kernel_size=(3, 3),
padding=(1, 1))
self.pixel_shuffle = PixelShuffle(up_scale)
self.p_relu = PReLU()
def dummy_test_space_to_depth():
f1 = SpaceToDim(2, dim=1)
ff = PixelShuffle(2)
x = Image.open('data/set5_x2/img_001_SRF_2_LR.png')
g = torchvision.transforms.ToTensor()
h = torchvision.transforms.ToPILImage()
z = f1(g(x).unsqueeze(0))
y = h(ff(z)[0])
assert np.all(np.array(x) == np.array(y))
def __init__(self, conv, scale, numFeatures, kernelSize, bias=True):
moduleList = []
if (scale & (scale - 1)) == 0: # Is scale = 2^n?
for i in range(int(math.log(scale, 2))):
moduleList.append(
conv(numFeatures, 4 * numFeatures, kernelSize, bias))
moduleList.append(PixelShuffle(2))
# if act: m.append(act())
elif scale == 3:
moduleList.append(
conv(numFeatures, 9 * numFeatures, kernelSize, bias))
moduleList.append(PixelShuffle(3))
# if act: m.append(act())
else:
raise NotImplementedError
super(Upsampler, self).__init__(*moduleList)
def __init__(self, upscale_factor):
super(NetXY, self).__init__()
self.conv1 = Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
self.conv2 = Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
# Here the 1 multiplicator can be 3 if the x and y outputs are wanted
self.conv3 = Conv2d(32, 1 * (upscale_factor**2), (3, 3), (1, 1),
(1, 1))
self.pixel_shuffle = PixelShuffle(upscale_factor)
def __init__(self, num_blocks=1):
super().__init__()
print("Initialized generator network..")
self.conv1 = Conv2d(1, 64, kernel_size=9, padding=4)
self.prelu = PReLU()
self.layers = self._get_residual_blocks(num_blocks)
self.conv2 = Conv2d(64, 64, kernel_size=3, padding=1)
self.bn2 = BatchNorm2d(64)
self.conv3 = Conv2d(64, 256, kernel_size=3, padding=1)
self.pxshuffle = PixelShuffle(upscale_factor=2) # up-sampling
# used in original SR-GAN paper, only for 4x up-sampling
# self.conv4 = Conv2d(256, 256, kernel_size=3, padding=1)
self.conv5 = Conv2d(64, 1, kernel_size=9, padding=4)
def __init__(self, input_channels: int, stride: int, mask_channels: int):
super().__init__()
self.expand = nn.Conv2d(input_channels, input_channels, kernel_size=1)
self.up = PixelShuffle(upscale_factor=stride)
mid_channels = input_channels // (2**stride)
self.conv1 = nn.Conv2d(mid_channels,
mid_channels,
kernel_size=3,
padding=1,
bias=False)
self.bn1 = nn.BatchNorm2d(mid_channels)
self.act1 = nn.ReLU(True)
self.conv2 = nn.Conv2d(mid_channels,
mid_channels,
kernel_size=3,
padding=1,
bias=False)
self.bn2 = nn.BatchNorm2d(mid_channels)
self.act2 = nn.ReLU(True)
self.final = nn.Conv2d(mid_channels, mask_channels, kernel_size=1)
def __init__(self):
super(UNet, self).__init__()
self.model_name = 'UNet'
# head block
self.head = Sequential(Conv2d(4, 32, 3, padding=1),
LeakyReLU(0.2, inplace=True),
Conv2d(32, 32, 3, padding=1),
LeakyReLU(0.2, inplace=True))
# block 1-4
self.d1 = downBlock(32, 64)
self.d2 = downBlock(64, 128)
self.d3 = downBlock(128, 256)
# bottom block
self.bottom = Sequential(MaxPool2d(2, 2), Conv2d(256,
512,
3,
padding=1),
LeakyReLU(0.2, inplace=True),
Conv2d(512, 512, 3, padding=1),
LeakyReLU(0.2, inplace=True),
ConvTranspose2d(512, 256, 2, stride=2))
# blcok 5-8
self.u1 = upBlock(512, 128)
self.u2 = upBlock(256, 64)
self.u3 = upBlock(128, 32)
# final block
self.final = Sequential(Conv2d(64, 32, 3, padding=1),
LeakyReLU(0.2, inplace=True),
Conv2d(32, 32, 3, padding=1),
LeakyReLU(0.2, inplace=True),
Conv2d(32, 12, 1), PixelShuffle(2))
def __init__(self):
super().__init__()
self.convnet = UNet(in_channels=4, out_channels=12)
self.upscaling = PixelShuffle(upscale_factor=2)
def __init__(self, in_channels, up_scale):
super(UpsampleBLock, self).__init__()
self.conv = Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)
self.pixel_shuffle = PixelShuffle(up_scale)
self.prelu = PReLU()