From 67bd678763ce5de5fa1804a630665d787c0369b9 Mon Sep 17 00:00:00 2001 From: protsenkovi Date: Thu, 11 Jul 2024 06:22:30 +0400 Subject: [PATCH] updates --- src/common/data.py | 2 + src/common/layers.py | 33 ++- src/common/losses.py | 28 +- src/models/sdynet.py | 368 +++++++++++++++++++++++++- src/models/srnet.py | 606 ++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 1019 insertions(+), 18 deletions(-) diff --git a/src/common/data.py b/src/common/data.py index 789e2ca..9298f78 100644 --- a/src/common/data.py +++ b/src/common/data.py @@ -13,6 +13,8 @@ from common.utils import pil2numpy image_extensions = ['.jpg', '.png'] 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]) cache_path = Path(images_dir_path).parent / f"{Path(images_dir_path).stem}_{color_model}_cache.npy" cache_path = cache_path.resolve() diff --git a/src/common/layers.py b/src/common/layers.py index b32c516..7088d0a 100644 --- a/src/common/layers.py +++ b/src/common/layers.py @@ -526,4 +526,35 @@ class ComplexGaborLayer2D(nn.Module): arg = scale_x.abs().square() + scale_y.abs().square() gauss_term = torch.exp(-self.scale_0*self.scale_0*arg) - return freq_term*gauss_term \ No newline at end of file + 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 \ No newline at end of file diff --git a/src/common/losses.py b/src/common/losses.py index 00ac9c6..d1d0765 100644 --- a/src/common/losses.py +++ b/src/common/losses.py @@ -401,7 +401,7 @@ class SSIM(torch.nn.Module): size_average: bool = True, win_size: int = 11, win_sigma: float = 1.5, - channel: int = 3, + # channel: int = 3, spatial_dims: int = 2, K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), nonnegative_ssim: bool = False, @@ -419,13 +419,16 @@ class SSIM(torch.nn.Module): super(SSIM, self).__init__() 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.data_range = data_range self.K = K self.nonnegative_ssim = nonnegative_ssim 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( X, Y, @@ -507,7 +510,6 @@ def get_outnorm(x:torch.Tensor, out_norm:str='') -> torch.Tensor: return norm - class CharbonnierLoss(nn.Module): """Charbonnier Loss (L1)""" def __init__(self, eps=1e-6, out_norm:str='bci'): @@ -518,4 +520,22 @@ class CharbonnierLoss(nn.Module): def forward(self, x, y): norm = get_outnorm(x, self.out_norm) loss = torch.sum(torch.sqrt((x - y).pow(2) + self.eps**2)) - return loss*norm \ No newline at end of file + 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 \ No newline at end of file diff --git a/src/models/sdynet.py b/src/models/sdynet.py index 14e133e..6ed3d6a 100644 --- a/src/models/sdynet.py +++ b/src/models/sdynet.py @@ -64,15 +64,15 @@ class SDYNetx2(SRNetBase): b,c,h,w = x.shape x = x.reshape(b*c, 1, h, w) 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, 1, 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_S, self.stage1_S) + output += self.forward_stage(x, self._extract_pattern_D, self.stage1_D) + output += self.forward_stage(x, self._extract_pattern_Y, self.stage1_Y) output /= 3 x = output 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.scale, 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_S, self.stage2_S) + output += self.forward_stage(x, self._extract_pattern_D, self.stage2_D) + output += self.forward_stage(x, self._extract_pattern_Y, self.stage2_Y) output /= 3 x = output 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) for rotations_count in range(4): 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, 1, 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_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.scale, 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.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_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) @@ -291,6 +291,59 @@ class SDYNetR90x2(SRNetBase): return F.mse_loss(pred/255, target/255) 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): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYEHONetR90x1, self).__init__() @@ -789,6 +842,297 @@ class SDYMixNetx1v7(SRNetBase): 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 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) diff --git a/src/models/srnet.py b/src/models/srnet.py index 41579fd..a55d3eb 100644 --- a/src/models/srnet.py +++ b/src/models/srnet.py @@ -33,6 +33,510 @@ class SRNet(SRNetBase): lut_model = srlut.SRLut.init_from_numpy(stage_lut) 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): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): @@ -626,4 +1130,104 @@ class SRMsbLsb4R90Net(SRNetBase): return x def get_lut_model(self, quantization_interval=16, batch_size=2**10): - raise NotImplementedError \ No newline at end of file + 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 \ No newline at end of file