main
protsenkovi 6 months ago
parent 274abed989
commit 67bd678763

@ -13,6 +13,8 @@ from common.utils import pil2numpy
image_extensions = ['.jpg', '.png'] image_extensions = ['.jpg', '.png']
def load_images_cached(images_dir_path, color_model, reset_cache): def load_images_cached(images_dir_path, color_model, reset_cache):
if not Path(images_dir_path).exists():
raise Exception(f"{images_dir_path} does not exist.")
image_paths = sorted([f for f in Path(images_dir_path).glob("*") if f.suffix.lower() in image_extensions]) image_paths = sorted([f for f in Path(images_dir_path).glob("*") if f.suffix.lower() in image_extensions])
cache_path = Path(images_dir_path).parent / f"{Path(images_dir_path).stem}_{color_model}_cache.npy" cache_path = Path(images_dir_path).parent / f"{Path(images_dir_path).stem}_{color_model}_cache.npy"
cache_path = cache_path.resolve() cache_path = cache_path.resolve()

@ -527,3 +527,34 @@ class ComplexGaborLayer2D(nn.Module):
gauss_term = torch.exp(-self.scale_0*self.scale_0*arg) gauss_term = torch.exp(-self.scale_0*self.scale_0*arg)
return freq_term*gauss_term return freq_term*gauss_term
class UpscaleBlockGabor(nn.Module):
def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255):
super(UpscaleBlockGabor, self).__init__()
assert layers_count > 0
self.upscale_factor = upscale_factor
self.hidden_dim = hidden_dim
self.embed = ComplexGaborLayer2D(in_features=in_features, out_features=hidden_dim, is_first=True)
self.linear_projections = []
for i in range(layers_count):
self.linear_projections.append(ComplexGaborLayer2D(in_features=hidden_dim, out_features=hidden_dim, is_first=False))
self.linear_projections = nn.ModuleList(self.linear_projections)
self.project_channels = nn.Linear(in_features=hidden_dim, out_features=upscale_factor * upscale_factor, bias=True)
self.in_bias = self.in_scale = input_max_value/2
self.out_bias = self.out_scale = output_max_value/2
self.layer_norm = nn.LayerNorm(hidden_dim) # To avoid gradient vanishing caused by tanh
def forward(self, x):
x = (x-self.in_bias)/self.in_scale
x = self.embed(x)
for linear_projection in self.linear_projections:
x = linear_projection(x)
x = x.real
x = self.project_channels(x)
x = torch.tanh(x)
x = x*self.out_scale + self.out_bias
return x

@ -401,7 +401,7 @@ class SSIM(torch.nn.Module):
size_average: bool = True, size_average: bool = True,
win_size: int = 11, win_size: int = 11,
win_sigma: float = 1.5, win_sigma: float = 1.5,
channel: int = 3, # channel: int = 3,
spatial_dims: int = 2, spatial_dims: int = 2,
K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), K: Union[Tuple[float, float], List[float]] = (0.01, 0.03),
nonnegative_ssim: bool = False, nonnegative_ssim: bool = False,
@ -419,13 +419,16 @@ class SSIM(torch.nn.Module):
super(SSIM, self).__init__() super(SSIM, self).__init__()
self.win_size = win_size self.win_size = win_size
self.win = _fspecial_gauss_1d(win_size, win_sigma).repeat([channel, 1] + [1] * spatial_dims) self.win_sigma = win_sigma
self.win = None
self.size_average = size_average self.size_average = size_average
self.data_range = data_range self.data_range = data_range
self.K = K self.K = K
self.nonnegative_ssim = nonnegative_ssim self.nonnegative_ssim = nonnegative_ssim
def forward(self, X: Tensor, Y: Tensor) -> Tensor: def forward(self, X: Tensor, Y: Tensor) -> Tensor:
if self.win is None:
self.win = _fspecial_gauss_1d(self.win_size, self.win_sigma).repeat([X.shape[1], 1] + [1] * len(X.shape[2:]))
return ssim( return ssim(
X, X,
Y, Y,
@ -507,7 +510,6 @@ def get_outnorm(x:torch.Tensor, out_norm:str='') -> torch.Tensor:
return norm return norm
class CharbonnierLoss(nn.Module): class CharbonnierLoss(nn.Module):
"""Charbonnier Loss (L1)""" """Charbonnier Loss (L1)"""
def __init__(self, eps=1e-6, out_norm:str='bci'): def __init__(self, eps=1e-6, out_norm:str='bci'):
@ -519,3 +521,21 @@ class CharbonnierLoss(nn.Module):
norm = get_outnorm(x, self.out_norm) norm = get_outnorm(x, self.out_norm)
loss = torch.sum(torch.sqrt((x - y).pow(2) + self.eps**2)) loss = torch.sum(torch.sqrt((x - y).pow(2) + self.eps**2))
return loss*norm return loss*norm
def huber(input, target, delta=0.01, reduce=True):
abs_error = torch.abs(input - target)
quadratic = torch.clamp(abs_error, max=delta)
# The following expression is the same in value as
# tf.maximum(abs_error - delta, 0), but importantly the gradient for the
# expression when abs_error == delta is 0 (for tf.maximum it would be 1).
# This is necessary to avoid doubling the gradient, since there is already a
# nonzero contribution to the gradient from the quadratic term.
linear = (abs_error - quadratic)
losses = 0.5 * torch.pow(quadratic, 2) + delta * linear
if reduce:
return torch.mean(losses)
else:
return losses

@ -64,15 +64,15 @@ class SDYNetx2(SRNetBase):
b,c,h,w = x.shape b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w) x = x.reshape(b*c, 1, h, w)
output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device) output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage1_S) output += self.forward_stage(x, self._extract_pattern_S, self.stage1_S)
output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage1_D) output += self.forward_stage(x, self._extract_pattern_D, self.stage1_D)
output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y) output += self.forward_stage(x, self._extract_pattern_Y, self.stage1_Y)
output /= 3 output /= 3
x = output x = output
output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S) output += self.forward_stage(x, self._extract_pattern_S, self.stage2_S)
output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D) output += self.forward_stage(x, self._extract_pattern_D, self.stage2_D)
output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y) output += self.forward_stage(x, self._extract_pattern_Y, self.stage2_Y)
output /= 3 output /= 3
x = output x = output
x = x.reshape(b, c, h*self.scale, w*self.scale) x = x.reshape(b, c, h*self.scale, w*self.scale)
@ -260,17 +260,17 @@ class SDYNetR90x2(SRNetBase):
output_1 = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device) output_1 = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
for rotations_count in range(4): for rotations_count in range(4):
rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1]) rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[-2, -1]) output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1]) output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1]) output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])
output_1 /= 4*3 output_1 /= 4*3
x = output_1 x = output_1
output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
for rotations_count in range(4): for rotations_count in range(4):
rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1]) rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage2_S), k=-rotations_count, dims=[-2, -1]) output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2_S), k=-rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1]) output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1]) output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1])
output_2 /= 4*3 output_2 /= 4*3
x = output_2 x = output_2
x = x.view(b, c, h*self.scale, w*self.scale) x = x.view(b, c, h*self.scale, w*self.scale)
@ -291,6 +291,59 @@ class SDYNetR90x2(SRNetBase):
return F.mse_loss(pred/255, target/255) return F.mse_loss(pred/255, target/255)
return loss_fn return loss_fn
class SDYNetR90x2Inv(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYNetR90x2Inv, self).__init__()
self.scale = scale
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_D = layers.PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_Y = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
self.stage1_S = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
self.stage1_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
self.stage1_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
self.stage2_S = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.view(b*c, 1, h, w)
output_1 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
for rotations_count in range(4):
rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1])
output_1 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])
output_1 /= 4*3
x = output_1
output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
for rotations_count in range(4):
rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2_S), k=-rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1])
output_2 += torch.rot90(self.forward_stage(rotated, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1])
output_2 /= 4*3
x = output_2
x = x.view(b, c, h*self.scale, w*self.scale)
return x
def get_lut_model(self, quantization_interval=16, batch_size=2**10):
stage1_S = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
stage1_D = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size)
stage1_Y = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size)
stage2_S = lut.transfer_2x2_input_SxS_output(self.stage2_S, quantization_interval=quantization_interval, batch_size=batch_size)
stage2_D = lut.transfer_2x2_input_SxS_output(self.stage2_D, quantization_interval=quantization_interval, batch_size=batch_size)
stage2_Y = lut.transfer_2x2_input_SxS_output(self.stage2_Y, quantization_interval=quantization_interval, batch_size=batch_size)
lut_model = sdylut.SDYLutR90x2.init_from_numpy(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y)
return lut_model
def get_loss_fn(self):
def loss_fn(pred, target):
return F.mse_loss(pred/255, target/255)
return loss_fn
class SDYEHONetR90x1(SRNetBase): class SDYEHONetR90x1(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYEHONetR90x1, self).__init__() super(SDYEHONetR90x1, self).__init__()
@ -793,3 +846,294 @@ class SDYMixNetx1v7(SRNetBase):
# return F.mse_loss(pred/255, target/255)# + ssim_loss # return F.mse_loss(pred/255, target/255)# + ssim_loss
return ssim_loss(pred, target) + l1_loss(pred, target) return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn return loss_fn
class SDYMixNetx1v8(SRNetBase):
"""
22
12 23 32 21
11 13 33 31
10 14 34 30
01 03 43 41
00 04 44 40
"""
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYMixNetx1v8, self).__init__()
self.scale = scale
self._extract_pattern_1 = layers.PercievePattern(receptive_field_idxes= [[3,3]], center=[3,3], window_size=7)
self._extract_pattern_2 = layers.PercievePattern(receptive_field_idxes= [[2,3],[3,4],[4,3],[3,2]], center=[3,3], window_size=7)
self._extract_pattern_3 = layers.PercievePattern(receptive_field_idxes= [[2,2],[2,4],[4,4],[4,2]], center=[3,3], window_size=7)
self._extract_pattern_4 = layers.PercievePattern(receptive_field_idxes= [[1,2],[1,4],[5,4],[5,2]], center=[3,3], window_size=7)
self._extract_pattern_5 = layers.PercievePattern(receptive_field_idxes= [[1,1],[1,5],[5,5],[5,1]], center=[3,3], window_size=7)
self._extract_pattern_6 = layers.PercievePattern(receptive_field_idxes= [[2,1],[2,5],[4,5],[4,1]], center=[3,3], window_size=7)
self._extract_pattern_7 = layers.PercievePattern(receptive_field_idxes= [[3,1],[1,3],[3,5],[5,3]], center=[3,3], window_size=7)
self._extract_pattern_8 = layers.PercievePattern(receptive_field_idxes= [[0,0],[0,6],[6,6],[6,0]], center=[3,3], window_size=7)
self._extract_pattern_9 = layers.PercievePattern(receptive_field_idxes= [[0,3],[3,6],[6,3],[3,0]], center=[3,3], window_size=7)
self._extract_pattern_10 = layers.PercievePattern(receptive_field_idxes=[[2,0],[1,0],[0,1],[0,2]], center=[3,3], window_size=7)
self._extract_pattern_11 = layers.PercievePattern(receptive_field_idxes=[[0,4],[0,5],[1,6],[2,6]], center=[3,3], window_size=7)
self._extract_pattern_12 = layers.PercievePattern(receptive_field_idxes=[[4,6],[5,6],[6,5],[6,4]], center=[3,3], window_size=7)
self._extract_pattern_13 = layers.PercievePattern(receptive_field_idxes=[[6,2],[6,1],[5,0],[4,0]], center=[3,3], window_size=7)
self.stage1_1 = layers.UpscaleBlock(in_features=1, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_2 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_3 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_4 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_5 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_6 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_7 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_8 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_9 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_10 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_11 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_12 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self.stage1_13 = layers.UpscaleBlock(in_features=4, hidden_dim=32, layers_count=2, upscale_factor=1)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=13)
self.stage1_Mix = layers.UpscaleBlock(in_features=13, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_1, self.stage1_1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_2, self.stage1_2)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_3, self.stage1_3)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_4, self.stage1_4)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_5, self.stage1_5)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_6, self.stage1_6)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_7, self.stage1_7)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_8, self.stage1_8)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_9, self.stage1_9)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_10, self.stage1_10)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_11, self.stage1_11)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_12, self.stage1_12)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_13, self.stage1_13)], dim=1)
output = self.forward_stage(output, self._extract_pattern_mix, self.stage1_Mix)
x = output
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
# return F.mse_loss(pred/255, target/255)# + ssim_loss
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SDYMixNetx1v9(SRNetBase):
"""
22
12 23 32 21
11 13 33 31
10 14 34 30
01 03 43 41
00 04 44 40
"""
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYMixNetx1v9, self).__init__()
self.scale = scale
self._extract_pattern_1 = layers.PercievePattern(receptive_field_idxes= [[3,3]], center=[3,3], window_size=7)
self._extract_pattern_2 = layers.PercievePattern(receptive_field_idxes= [[2,3],[3,4],[4,3],[3,2]], center=[3,3], window_size=7)
self._extract_pattern_3 = layers.PercievePattern(receptive_field_idxes= [[2,2],[2,4],[4,4],[4,2]], center=[3,3], window_size=7)
self._extract_pattern_4 = layers.PercievePattern(receptive_field_idxes= [[1,2],[1,4],[5,4],[5,2]], center=[3,3], window_size=7)
self._extract_pattern_5 = layers.PercievePattern(receptive_field_idxes= [[1,1],[1,5],[5,5],[5,1]], center=[3,3], window_size=7)
self._extract_pattern_6 = layers.PercievePattern(receptive_field_idxes= [[2,1],[2,5],[4,5],[4,1]], center=[3,3], window_size=7)
self._extract_pattern_7 = layers.PercievePattern(receptive_field_idxes= [[3,1],[1,3],[3,5],[5,3]], center=[3,3], window_size=7)
self._extract_pattern_8 = layers.PercievePattern(receptive_field_idxes= [[0,0],[0,6],[6,6],[6,0]], center=[3,3], window_size=7)
self._extract_pattern_9 = layers.PercievePattern(receptive_field_idxes= [[0,3],[3,6],[6,3],[3,0]], center=[3,3], window_size=7)
self._extract_pattern_10 = layers.PercievePattern(receptive_field_idxes=[[2,0],[1,0],[0,1],[0,2]], center=[3,3], window_size=7)
self._extract_pattern_11 = layers.PercievePattern(receptive_field_idxes=[[0,4],[0,5],[1,6],[2,6]], center=[3,3], window_size=7)
self._extract_pattern_12 = layers.PercievePattern(receptive_field_idxes=[[4,6],[5,6],[6,5],[6,4]], center=[3,3], window_size=7)
self._extract_pattern_13 = layers.PercievePattern(receptive_field_idxes=[[6,2],[6,1],[5,0],[4,0]], center=[3,3], window_size=7)
self.stage1_1 = layers.UpscaleBlock(in_features=1, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_4 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_5 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_6 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_7 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_8 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_9 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_10 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_11 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_12 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_13 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=13)
self.stage1_Mix = layers.UpscaleBlockChebyKAN(in_features=13, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_1, self.stage1_1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_2, self.stage1_2)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_3, self.stage1_3)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_4, self.stage1_4)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_5, self.stage1_5)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_6, self.stage1_6)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_7, self.stage1_7)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_8, self.stage1_8)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_9, self.stage1_9)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_10, self.stage1_10)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_11, self.stage1_11)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_12, self.stage1_12)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_13, self.stage1_13)], dim=1)
x = self.forward_stage(output, self._extract_pattern_mix, self.stage1_Mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
# return F.mse_loss(pred/255, target/255)# + ssim_loss
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SDYMixNetx1v10(SRNetBase):
"""
22
12 23 32 21
11 13 33 31
10 14 34 30
01 03 43 41
00 04 44 40
"""
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYMixNetx1v10, self).__init__()
self.scale = scale
self._extract_pattern_1 = layers.PercievePattern(receptive_field_idxes= [[3,3]], center=[3,3], window_size=7)
self._extract_pattern_2 = layers.PercievePattern(receptive_field_idxes= [[2,3],[3,4],[4,3],[3,2]], center=[3,3], window_size=7)
self._extract_pattern_3 = layers.PercievePattern(receptive_field_idxes= [[2,2],[2,4],[4,4],[4,2]], center=[3,3], window_size=7)
self._extract_pattern_4 = layers.PercievePattern(receptive_field_idxes= [[1,2],[1,4],[5,4],[5,2]], center=[3,3], window_size=7)
self._extract_pattern_5 = layers.PercievePattern(receptive_field_idxes= [[1,1],[1,5],[5,5],[5,1]], center=[3,3], window_size=7)
self._extract_pattern_6 = layers.PercievePattern(receptive_field_idxes= [[2,1],[2,5],[4,5],[4,1]], center=[3,3], window_size=7)
self._extract_pattern_7 = layers.PercievePattern(receptive_field_idxes= [[3,1],[1,3],[3,5],[5,3]], center=[3,3], window_size=7)
self._extract_pattern_8 = layers.PercievePattern(receptive_field_idxes= [[0,0],[0,6],[6,6],[6,0]], center=[3,3], window_size=7)
self._extract_pattern_9 = layers.PercievePattern(receptive_field_idxes= [[0,3],[3,6],[6,3],[3,0]], center=[3,3], window_size=7)
self._extract_pattern_10 = layers.PercievePattern(receptive_field_idxes=[[2,0],[1,0],[0,1],[0,2]], center=[3,3], window_size=7)
self._extract_pattern_11 = layers.PercievePattern(receptive_field_idxes=[[0,4],[0,5],[1,6],[2,6]], center=[3,3], window_size=7)
self._extract_pattern_12 = layers.PercievePattern(receptive_field_idxes=[[4,6],[5,6],[6,5],[6,4]], center=[3,3], window_size=7)
self._extract_pattern_13 = layers.PercievePattern(receptive_field_idxes=[[6,2],[6,1],[5,0],[4,0]], center=[3,3], window_size=7)
self.stage1_1 = layers.UpscaleBlock(in_features=1, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_4 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_5 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_6 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_7 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_8 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_9 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_10 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_11 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_12 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self.stage1_13 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=1)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=13)
self.stage1_Mix = layers.UpscaleBlockGabor(in_features=13, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_1, self.stage1_1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_2, self.stage1_2)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_3, self.stage1_3)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_4, self.stage1_4)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_5, self.stage1_5)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_6, self.stage1_6)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_7, self.stage1_7)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_8, self.stage1_8)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_9, self.stage1_9)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_10, self.stage1_10)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_11, self.stage1_11)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_12, self.stage1_12)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_13, self.stage1_13)], dim=1)
x = self.forward_stage(output, self._extract_pattern_mix, self.stage1_Mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
# return F.mse_loss(pred/255, target/255)# + ssim_loss
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SDYMixNetx1v11(SRNetBase):
"""
22
12 23 32 21
11 13 33 31
10 14 34 30
01 03 43 41
00 04 44 40
"""
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SDYMixNetx1v11, self).__init__()
self.scale = scale
self._extract_pattern_1 = layers.PercievePattern(receptive_field_idxes= [[3,3]], center=[3,3], window_size=7)
self._extract_pattern_2 = layers.PercievePattern(receptive_field_idxes= [[2,3],[3,4],[4,3],[3,2]], center=[3,3], window_size=7)
self._extract_pattern_3 = layers.PercievePattern(receptive_field_idxes= [[2,2],[2,4],[4,4],[4,2]], center=[3,3], window_size=7)
self._extract_pattern_4 = layers.PercievePattern(receptive_field_idxes= [[1,2],[1,4],[5,4],[5,2]], center=[3,3], window_size=7)
self._extract_pattern_5 = layers.PercievePattern(receptive_field_idxes= [[1,1],[1,5],[5,5],[5,1]], center=[3,3], window_size=7)
self._extract_pattern_6 = layers.PercievePattern(receptive_field_idxes= [[2,1],[2,5],[4,5],[4,1]], center=[3,3], window_size=7)
self._extract_pattern_7 = layers.PercievePattern(receptive_field_idxes= [[3,1],[1,3],[3,5],[5,3]], center=[3,3], window_size=7)
self._extract_pattern_8 = layers.PercievePattern(receptive_field_idxes= [[0,0],[0,6],[6,6],[6,0]], center=[3,3], window_size=7)
self._extract_pattern_9 = layers.PercievePattern(receptive_field_idxes= [[0,3],[3,6],[6,3],[3,0]], center=[3,3], window_size=7)
self._extract_pattern_10 = layers.PercievePattern(receptive_field_idxes=[[2,0],[1,0],[0,1],[0,2]], center=[3,3], window_size=7)
self._extract_pattern_11 = layers.PercievePattern(receptive_field_idxes=[[0,4],[0,5],[1,6],[2,6]], center=[3,3], window_size=7)
self._extract_pattern_12 = layers.PercievePattern(receptive_field_idxes=[[4,6],[5,6],[6,5],[6,4]], center=[3,3], window_size=7)
self._extract_pattern_13 = layers.PercievePattern(receptive_field_idxes=[[6,2],[6,1],[5,0],[4,0]], center=[3,3], window_size=7)
self.stage1_1 = layers.UpscaleBlock(in_features=1, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_4 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_5 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_6 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_7 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_8 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_9 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_10 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_11 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_12 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self.stage1_13 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=1, upscale_factor=scale)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=13)
self.stage1_Mix = layers.UpscaleBlockGabor(in_features=13, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_1, self.stage1_1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_2, self.stage1_2)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_3, self.stage1_3)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_4, self.stage1_4)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_5, self.stage1_5)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_6, self.stage1_6)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_7, self.stage1_7)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_8, self.stage1_8)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_9, self.stage1_9)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_10, self.stage1_10)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_11, self.stage1_11)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_12, self.stage1_12)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_13, self.stage1_13)], dim=1)
x = self.forward_stage(output, self._extract_pattern_mix, self.stage1_Mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
# return F.mse_loss(pred/255, target/255)# + ssim_loss
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn

@ -33,6 +33,510 @@ class SRNet(SRNetBase):
lut_model = srlut.SRLut.init_from_numpy(stage_lut) lut_model = srlut.SRLut.init_from_numpy(stage_lut)
return lut_model return lut_model
class SRNetMixer(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixer, self).__init__()
self.scale = scale
self.stage1_1 = layers.UpscaleBlock(in_features=1, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_4 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_5 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_6 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_7 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_8 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_9 = layers.UpscaleBlock(in_features=1, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_10 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_11 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_12 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage1_13 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self._extract_pattern_S1 = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1)
self._extract_pattern_S2 = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,2],[2,0],[2,2]], center=[1,1], window_size=3)
self._extract_pattern_S3 = layers.PercievePattern(receptive_field_idxes=[[0,1],[1,2],[2,1],[1,0]], center=[1,1], window_size=3)
self._extract_pattern_S4 = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,4],[4,4],[4,0]], center=[1,1], window_size=5)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S1, self.stage1_1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S2, self.stage1_2)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S3, self.stage1_3)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S4, self.stage1_4)], dim=1)
x = output
output = self.forward_stage(x, self._extract_pattern_mix, self.stage1_5)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_mix, self.stage1_6)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_mix, self.stage1_7)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_mix, self.stage1_8)], dim=1)
x = output
output = self.forward_stage(x, self._extract_pattern_S1, self.stage1_9)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S2, self.stage1_10)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S3, self.stage1_11)], dim=1)
output = torch.cat([output, self.forward_stage(x, self._extract_pattern_S4, self.stage1_12)], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage1_13)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1) for _ in range(4)])
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v2(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v2, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale) for _ in range(4)])
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v3(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v3, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale//2) for _ in range(4)])
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale//2)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v5(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v5, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale) for _ in range(4)])
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1) for _ in range(4)])
self.stage_mix2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix)
output = self.forward_stage(x, self._extract_pattern_S, self.stage2[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2[rotations_count]), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix2)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v9(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v9, self).__init__()
self.scale = scale
self.stage1 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage_mix2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix)
output = self.forward_stage(x, self._extract_pattern_S, self.stage2)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix2)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v10(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v10, self).__init__()
self.scale = scale
self.stage1S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1D = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1Y = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage_mix2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_D = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,2],[2,0],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_Y = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def stage(self, x, pattern, stage_net, mix_pattern, stage_mix_net):
tmp_output = self.forward_stage(x, pattern, stage_net)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, pattern, stage_net), k=-rotations_count, dims=[2, 3])], dim=1)
output = self.forward_stage(tmp_output, mix_pattern, stage_mix_net)
return output
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.stage(x, self._extract_pattern_S, self.stage1S, self._extract_pattern_mix, self.stage_mix)
output += self.stage(x, self._extract_pattern_D, self.stage1D, self._extract_pattern_mix, self.stage_mix)
output += self.stage(x, self._extract_pattern_Y, self.stage1Y, self._extract_pattern_mix, self.stage_mix)
x = output / 3
x = x.clamp(0, 255)
x = self.stage(x, self._extract_pattern_S, self.stage2S, self._extract_pattern_mix, self.stage_mix2)
x = x.clamp(0, 255)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v10Lsb(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v10Lsb, self).__init__()
self.scale = scale
self.stage1S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1D = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1Y = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage2S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage_mix2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage1S_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=15, output_max_value=15)
self.stage1D_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=15, output_max_value=15)
self.stage1Y_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=15, output_max_value=15)
self.stage_mix_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1, input_max_value=15, output_max_value=255)
self.stage2S_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=15, output_max_value=15)
self.stage_mix2_lsb = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=15, output_max_value=255)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_D = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,2],[2,0],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_Y = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def stage(self, x, pattern, stage_net, mix_pattern, stage_mix_net):
tmp_output = self.forward_stage(x, pattern, stage_net)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, pattern, stage_net), k=-rotations_count, dims=[2, 3])], dim=1)
output = self.forward_stage(tmp_output, mix_pattern, stage_mix_net)
return output
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
lsb = x % 16
msb = x - lsb
output_msb = self.stage(msb, self._extract_pattern_S, self.stage1S, self._extract_pattern_mix, self.stage_mix)
output_msb += self.stage(msb, self._extract_pattern_D, self.stage1D, self._extract_pattern_mix, self.stage_mix)
output_msb += self.stage(msb, self._extract_pattern_Y, self.stage1Y, self._extract_pattern_mix, self.stage_mix)
output_lsb = self.stage(lsb, self._extract_pattern_S, self.stage1S_lsb, self._extract_pattern_mix, self.stage_mix_lsb)
output_lsb += self.stage(lsb, self._extract_pattern_D, self.stage1D_lsb, self._extract_pattern_mix, self.stage_mix_lsb)
output_lsb += self.stage(lsb, self._extract_pattern_Y, self.stage1Y_lsb, self._extract_pattern_mix, self.stage_mix_lsb)
x = output_msb / 3 + output_lsb / 3
x = x.clamp(0, 255)
lsb = x % 16
msb = x - lsb
output_msb = self.stage(msb, self._extract_pattern_S, self.stage2S, self._extract_pattern_mix, self.stage_mix2)
output_lsb = self.stage(lsb, self._extract_pattern_S, self.stage2S_lsb, self._extract_pattern_mix, self.stage_mix2_lsb)
x = output_msb + output_lsb
x = x.clamp(0, 255)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v4(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v4, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale) for _ in range(4)])
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])
x = output
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v6(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v6, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale) for _ in range(4)])
self.stage2 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1) for _ in range(4)])
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])
x = output
output = self.forward_stage(x, self._extract_pattern_S, self.stage2[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2[rotations_count]), k=-rotations_count, dims=[2, 3])
x = output
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(channel=1, data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
class SRNetMixerR90v7(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v7, self).__init__()
self.scale = scale
self.stage1 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale) for _ in range(4)])
self.stage2 = nn.ModuleList([layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1) for _ in range(4)])
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
output = self.forward_stage(x, self._extract_pattern_S, self.stage1[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1[rotations_count]), k=-rotations_count, dims=[2, 3])
x = output
output = self.forward_stage(x, self._extract_pattern_S, self.stage2[0])
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output += torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2[rotations_count]), k=-rotations_count, dims=[2, 3])
x = output
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
class HDBNetv4(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4):
super(HDBNetv4, self).__init__()
assert scale == 4
self.scale = scale
self.rotations = rotations
self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=255, output_max_value=255)
self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=255, output_max_value=255)
self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=255, output_max_value=255)
self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=15, output_max_value=255)
self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale, input_max_value=15, output_max_value=255)
self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=255, output_max_value=255)
self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=255, output_max_value=255)
self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=255, output_max_value=255)
self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=15, output_max_value=255)
self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1, input_max_value=15, output_max_value=255)
self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
upsampled = nn.Upsample(scale_factor=self.scale, mode='nearest')(x)
lsb = x % 16
msb = x - lsb
output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
for rotations_count in range(self.rotations):
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
output_msbt = self.forward_stage(rotated_msb, self._extract_pattern_3H, self.stage1_3H) + \
self.forward_stage(rotated_msb, self._extract_pattern_3D, self.stage1_3D) + \
self.forward_stage(rotated_msb, self._extract_pattern_3B, self.stage1_3B)
output_lsbt = self.forward_stage(rotated_lsb, self._extract_pattern_2H, self.stage1_2H) + \
self.forward_stage(rotated_lsb, self._extract_pattern_2D, self.stage1_2D)
if not config is None and config.current_iter % config.display_step == 0:
config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
output_msb /= self.rotations*3
output_lsb /= self.rotations*2
output = output_msb + output_lsb
x = output.clamp(0, 255)
lsb = x % 16
msb = x - lsb
output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale,], dtype=x.dtype, device=x.device)
output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale,], dtype=x.dtype, device=x.device)
for rotations_count in range(self.rotations):
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
output_msbt = self.forward_stage(rotated_msb, self._extract_pattern_3H, self.stage2_3H) + \
self.forward_stage(rotated_msb, self._extract_pattern_3D, self.stage2_3D) + \
self.forward_stage(rotated_msb, self._extract_pattern_3B, self.stage2_3B)
output_lsbt = self.forward_stage(rotated_lsb, self._extract_pattern_2H, self.stage2_2H) + \
self.forward_stage(rotated_lsb, self._extract_pattern_2D, self.stage2_2D)
if not config is None and config.current_iter % config.display_step == 0:
config.writer.add_histogram('s2_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
config.writer.add_histogram('s2_output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
output_msb /= self.rotations*3
output_lsb /= self.rotations*2
output_msb -= 127.5
output_msb -= 127.5
output = output_msb + output_lsb + upsampled
x = output.clamp(0, 255)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
class SRNetGabor(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetGabor, self).__init__()
self.scale = scale
self.stage1_S = layers.UpscaleBlockGabor(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
x = self.forward_stage(x, self._extract_pattern_S, self.stage1_S)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_lut_model(self, quantization_interval=16, batch_size=2**10):
stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
lut_model = srlut.SRLut.init_from_numpy(stage_lut)
return lut_model
class SRNetChebyKan(SRNetBase): class SRNetChebyKan(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
@ -627,3 +1131,103 @@ class SRMsbLsb4R90Net(SRNetBase):
def get_lut_model(self, quantization_interval=16, batch_size=2**10): def get_lut_model(self, quantization_interval=16, batch_size=2**10):
raise NotImplementedError raise NotImplementedError
class SRNetMixerR90v11(SRNetBase):
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
super(SRNetMixerR90v11, self).__init__()
self.scale = scale
self.stage1S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1D = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage1Y = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage2S = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage_mix = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
self.stage_mix2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self.stage_mix3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=2, upscale_factor=1)
self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
self._extract_pattern_D = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,2],[2,0],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_Y = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1],[2,2]], center=[0,0], window_size=3)
self._extract_pattern_mix = layers.PercievePattern(receptive_field_idxes=[[0,0]], center=[0,0], window_size=1, channels=4)
self._extract_pattern_1 = layers.PercievePattern(receptive_field_idxes=[[3,0],[3,1],[3,2],[3,3]], center=[0,0], window_size=4)
self._extract_pattern_2 = layers.PercievePattern(receptive_field_idxes=[[0,3],[1,3],[2,3],[3,3]], center=[0,0], window_size=4)
self._extract_pattern_3 = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2],[3,2]], center=[0,0], window_size=4)
self._extract_pattern_4 = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,4],[4,0],[4,4]], center=[0,0], window_size=5)
self.stage_pattern_1 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_pattern_2 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_pattern_3 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
self.stage_pattern_4 = layers.UpscaleBlock(in_features=4, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
def forward(self, x, config=None):
b,c,h,w = x.shape
x = x.reshape(b*c, 1, h, w)
tmp_output = self.forward_stage(x, self._extract_pattern_S, self.stage1S)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage1S), k=-rotations_count, dims=[2, 3])], dim=1)
output = self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix)
tmp_output = self.forward_stage(x, self._extract_pattern_D, self.stage1D)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_D, self.stage1D), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix)
tmp_output = self.forward_stage(x, self._extract_pattern_Y, self.stage1Y)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_Y, self.stage1Y), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix)
########################################
tmp_output = self.forward_stage(x, self._extract_pattern_1, self.stage_pattern_1)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_1, self.stage_pattern_1), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix2)
tmp_output = self.forward_stage(x, self._extract_pattern_2, self.stage_pattern_2)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_2, self.stage_pattern_2), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix2)
tmp_output = self.forward_stage(x, self._extract_pattern_3, self.stage_pattern_3)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_3, self.stage_pattern_3), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix2)
tmp_output = self.forward_stage(x, self._extract_pattern_4, self.stage_pattern_4)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
tmp_output = torch.cat([tmp_output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_4, self.stage_pattern_4), k=-rotations_count, dims=[2, 3])], dim=1)
output += self.forward_stage(tmp_output, self._extract_pattern_mix, self.stage_mix2)
output /= 7
x = output
x = x.clamp(0, 255)
output = self.forward_stage(x, self._extract_pattern_S, self.stage2S)
for rotations_count in range(1,4):
rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
output = torch.cat([output, torch.rot90(self.forward_stage(rotated, self._extract_pattern_S, self.stage2S), k=-rotations_count, dims=[2, 3])], dim=1)
x = output
x = self.forward_stage(x, self._extract_pattern_mix, self.stage_mix3)
x = x.clamp(0, 255)
x = x.reshape(b, c, h*self.scale, w*self.scale)
return x
def get_loss_fn(self):
ssim_loss = losses.SSIM(data_range=255)
l1_loss = losses.CharbonnierLoss()
def loss_fn(pred, target):
return ssim_loss(pred, target) + l1_loss(pred, target)
return loss_fn
Loading…
Cancel
Save