From 221e3dc865f8d63c4e04302ea180396ad452d06a Mon Sep 17 00:00:00 2001 From: protsenkovi Date: Tue, 18 Jun 2024 15:57:10 +0400 Subject: [PATCH] v2 r90 anf flip msblsb models. --- src/models/srnet.py | 100 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/src/models/srnet.py b/src/models/srnet.py index dc3a03d..373e160 100644 --- a/src/models/srnet.py +++ b/src/models/srnet.py @@ -370,6 +370,106 @@ class SRMsbLsbR90Net(SRNetBase): raise NotImplementedError +class SRMsbLsbR90v2Net(SRNetBase): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SRMsbLsbR90v2Net, self).__init__() + self.scale = scale + self.hidden_dim = hidden_dim + self.layers_count = layers_count + + self.msb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + input_max_value=255, + output_max_value=255 + ) + self.lsb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + 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) + + def forward(self, x, config=None): + b,c,h,w = x.shape + x = x.reshape(b*c, 1, h, w) + output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) + for rotations_count in range(4): + rot_x = torch.rot90(x, k=rotations_count, dims=[2, 3]) + rotated_lsb = rot_x % 16 + rotated_msb = rot_x - rotated_lsb + output_msb = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, self.msb_fn) + output_lsb = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, self.lsb_fn) + if not config is None and config.current_iter % config.display_step == 0: + config.writer.add_histogram('output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter) + config.writer.add_histogram('output_msb', output_msb.detach().cpu().numpy(), config.current_iter) + output += torch.rot90(output_msb + output_lsb, k=-rotations_count, dims=[2, 3]) + output /= 4 + x = output + 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): + raise NotImplementedError + +class SRMsbLsbFlipNet(SRNetBase): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SRMsbLsbFlipNet, self).__init__() + self.scale = scale + self.hidden_dim = hidden_dim + self.layers_count = layers_count + + self.msb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + input_max_value=255, + output_max_value=255 + ) + self.lsb_fn = layers.UpscaleBlock( + in_features=4, + hidden_dim=hidden_dim, + layers_count=layers_count, + upscale_factor=self.scale, + 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.flip_functions = [ + lambda x: x, + lambda x: x[:,:,::-1,:], + lambda x: x[:,:,:,::-1], + lambda x: x[:,:,::-1,::-1], + ] + + def forward(self, x, config=None): + b,c,h,w = x.shape + x = x.reshape(b*c, 1, h, w) + output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) + for flip_f in self.flips_functions: + fliped_x = flip_f(x) + fliped_lsb = fliped_x % 16 + fliped_msb = fliped_x - fliped_lsb + output_msb = self.forward_stage(fliped_msb, self.scale, self._extract_pattern_S, self.msb_fn) + output_lsb = self.forward_stage(fliped_lsb, self.scale, self._extract_pattern_S, self.lsb_fn) + if not config is None and config.current_iter % config.display_step == 0: + config.writer.add_histogram('output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter) + config.writer.add_histogram('output_msb', output_msb.detach().cpu().numpy(), config.current_iter) + output += torch.rot90(output_msb + output_lsb, k=-rotations_count, dims=[2, 3]) + output /= 4 + x = output + 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): + raise NotImplementedError + class SRMsbLsb4R90Net(SRNetBase): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SRMsbLsb4R90Net, self).__init__()