|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.nn.functional as F
|
|
|
import numpy as np
|
|
|
|
|
|
class PercievePattern():
|
|
|
def __init__(self, receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2):
|
|
|
assert window_size >= (np.max(receptive_field_idxes)+1)
|
|
|
self.receptive_field_idxes = np.array(receptive_field_idxes)
|
|
|
self.window_size = window_size
|
|
|
self.center = center
|
|
|
self.receptive_field_idxes = [
|
|
|
self.receptive_field_idxes[0,0]*self.window_size + self.receptive_field_idxes[0,1],
|
|
|
self.receptive_field_idxes[1,0]*self.window_size + self.receptive_field_idxes[1,1],
|
|
|
self.receptive_field_idxes[2,0]*self.window_size + self.receptive_field_idxes[2,1],
|
|
|
self.receptive_field_idxes[3,0]*self.window_size + self.receptive_field_idxes[3,1],
|
|
|
]
|
|
|
|
|
|
def __call__(self, x):
|
|
|
b,c,h,w = x.shape
|
|
|
x = F.pad(x, pad=[self.center[0], self.window_size-self.center[0]-1,
|
|
|
self.center[1], self.window_size-self.center[1]-1], mode='replicate')
|
|
|
x = F.unfold(input=x, kernel_size=self.window_size)
|
|
|
x = torch.stack([
|
|
|
x[:,self.receptive_field_idxes[0],:],
|
|
|
x[:,self.receptive_field_idxes[1],:],
|
|
|
x[:,self.receptive_field_idxes[2],:],
|
|
|
x[:,self.receptive_field_idxes[3],:]
|
|
|
], 2)
|
|
|
x = x.reshape(x.shape[0]*x.shape[1], 1, 2, 2)
|
|
|
return x
|
|
|
|
|
|
# Huang G. et al. Densely connected convolutional networks //Proceedings of the IEEE conference on computer vision and pattern recognition. – 2017. – С. 4700-4708.
|
|
|
# https://ar5iv.labs.arxiv.org/html/1608.06993
|
|
|
# https://github.com/andreasveit/densenet-pytorch/blob/63152f4a40644b62717749536ed2e011c6e4d9ab/densenet.py#L40
|
|
|
class DenseConvUpscaleBlock(nn.Module):
|
|
|
def __init__(self, hidden_dim = 32, layers_count=5, upscale_factor=1):
|
|
|
super(DenseConvUpscaleBlock, self).__init__()
|
|
|
assert layers_count > 0
|
|
|
self.upscale_factor = upscale_factor
|
|
|
self.hidden_dim = hidden_dim
|
|
|
self.percieve = nn.Conv2d(1, hidden_dim, kernel_size=(2, 2), padding='valid', stride=1, dilation=1, bias=True)
|
|
|
self.convs = []
|
|
|
for i in range(layers_count):
|
|
|
self.convs.append(nn.Conv2d(in_channels = (i+1)*hidden_dim, out_channels = hidden_dim, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True))
|
|
|
self.convs = nn.ModuleList(self.convs)
|
|
|
|
|
|
for name, p in self.named_parameters():
|
|
|
if "weight" in name: nn.init.kaiming_normal_(p)
|
|
|
if "bias" in name: nn.init.constant_(p, 0)
|
|
|
|
|
|
self.project_channels = nn.Conv2d(in_channels = (layers_count+1)*hidden_dim, out_channels = upscale_factor * upscale_factor, kernel_size = 1, stride=1, padding=0, dilation=1, bias=True)
|
|
|
self.shuffle = nn.PixelShuffle(upscale_factor)
|
|
|
|
|
|
def forward(self, x):
|
|
|
x = (x-127.5)/127.5
|
|
|
x = torch.relu(self.percieve(x))
|
|
|
for conv in self.convs:
|
|
|
x = torch.cat([x, torch.relu(conv(x))], dim=1)
|
|
|
x = self.shuffle(self.project_channels(x))
|
|
|
x = torch.tanh(x)
|
|
|
x = round_func(x*127.5 + 127.5)
|
|
|
return x |