From b641401ed71ca1d96f79deceb01b8eea76da3cbe Mon Sep 17 00:00:00 2001 From: protsenkovi Date: Mon, 20 May 2024 11:50:14 +0400 Subject: [PATCH] bugfix, new model --- src/models/__init__.py | 2 ++ src/models/srlut.py | 47 ++++++++++++++++++++++++++++++++++ src/models/srnet.py | 57 ++++++++++++++++++++++++++++++++++++------ 3 files changed, 98 insertions(+), 8 deletions(-) diff --git a/src/models/__init__.py b/src/models/__init__.py index b1ba782..9001ba5 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -18,6 +18,8 @@ AVAILABLE_MODELS = { 'SDYLutx1': sdylut.SDYLutx1, 'SDYNetx2': sdynet.SDYNetx2, 'SDYLutx2': sdylut.SDYLutx2, + 'SRNetY': srnet.SRNetY, + 'SRLutY': srlut.SRLutY, # 'RCNetCentered_3x3': rcnet.RCNetCentered_3x3, 'RCLutCentered_3x3': rclut.RCLutCentered_3x3, # 'RCNetCentered_7x7': rcnet.RCNetCentered_7x7, 'RCLutCentered_7x7': rclut.RCLutCentered_7x7, # 'RCNetRot90_3x3': rcnet.RCNetRot90_3x3, 'RCLutRot90_3x3': rclut.RCLutRot90_3x3, diff --git a/src/models/srlut.py b/src/models/srlut.py index f8a591d..d5e2be9 100644 --- a/src/models/srlut.py +++ b/src/models/srlut.py @@ -49,7 +49,54 @@ class SRLut(nn.Module): def __repr__(self): return f"{self.__class__.__name__}\n lut size: {self.stage_lut.shape}" +class SRLutY(nn.Module): + def __init__( + self, + quantization_interval, + scale + ): + super(SRLutY, self).__init__() + self.scale = scale + self.quantization_interval = quantization_interval + self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + self.rgb_to_ycbcr = layers.RgbToYcbcr() + self.ycbcr_to_rgb = layers.YcbcrToRgb() + @staticmethod + def init_from_numpy( + stage_lut + ): + scale = int(stage_lut.shape[-1]) + quantization_interval = 256//(stage_lut.shape[0]-1) + lut_model = SRLutY(quantization_interval=quantization_interval, scale=scale) + lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).type(torch.float32)) + return lut_model + + def forward_stage(self, x, scale, percieve_pattern, lut): + b,c,h,w = x.shape + x = percieve_pattern(x) + x = select_index_4dlut_tetrahedral(index=x, lut=lut) + x = round_func(x) + x = x.reshape(b, c, h, w, scale, scale) + x = x.permute(0,1,2,4,3,5) + x = x.reshape(b, c, h*scale, w*scale) + return x + + def forward(self, x): + b,c,h,w = x.shape + x = self.rgb_to_ycbcr(x) + y = x[:,0:1,:,:] + cbcr = x[:,1:,:,:] + cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear') + + output = self.forward_stage(y, self.scale, self._extract_pattern_S, self.stage_lut) + output = torch.cat([output, cbcr_scaled], dim=1) + output = self.ycbcr_to_rgb(output).clamp(0, 255) + return output + + def __repr__(self): + return f"{self.__class__.__name__}\n lut size: {self.stage_lut.shape}" class SRLutR90(nn.Module): def __init__( diff --git a/src/models/srnet.py b/src/models/srnet.py index d759ef1..3623a6a 100644 --- a/src/models/srnet.py +++ b/src/models/srnet.py @@ -17,7 +17,7 @@ class SRNet(nn.Module): layers_count=layers_count, upscale_factor=self.scale ) - self._unfold_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + 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_stage(self, x, scale, percieve_pattern, stage): b,c,h,w = x.shape @@ -32,7 +32,7 @@ class SRNet(nn.Module): def forward(self, x): b,c,h,w = x.shape x = x.reshape(b*c, 1, h, w) - x = self.forward_stage(x, self.scale, self._unfold_pattern_S, self.stage1_S) + x = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S) x = x.reshape(b, c, h*self.scale, w*self.scale) return x @@ -41,6 +41,47 @@ class SRNet(nn.Module): lut_model = srlut.SRLut.init_from_numpy(stage_lut) return lut_model +class SRNetY(nn.Module): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SRNetY, self).__init__() + self.scale = scale + self.stage1_S = layers.UpscaleBlock( + 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.rgb_to_ycbcr = layers.RgbToYcbcr() + self.ycbcr_to_rgb = layers.YcbcrToRgb() + + def forward_stage(self, x, scale, percieve_pattern, stage): + b,c,h,w = x.shape + x = percieve_pattern(x) + x = stage(x) + x = round_func(x) + x = x.reshape(b, c, h, w, scale, scale) + x = x.permute(0,1,2,4,3,5) + x = x.reshape(b, c, h*scale, w*scale) + return x + + def forward(self, x): + b,c,h,w = x.shape + x = self.rgb_to_ycbcr(x) + y = x[:,0:1,:,:] + cbcr = x[:,1:,:,:] + cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear') + + x = y.view(b, 1, h, w) + output = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S) + output = torch.cat([output, cbcr_scaled], dim=1) + output = self.ycbcr_to_rgb(output).clamp(0, 255) + return output + + def get_lut_model(self, quantization_interval=16, batch_size=2**10): + stage_lut = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size) + lut_model = srlut.SRLutY.init_from_numpy(stage_lut) + return lut_model + class SRNetR90(nn.Module): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SRNetR90, self).__init__() @@ -50,7 +91,7 @@ class SRNetR90(nn.Module): layers_count=layers_count, upscale_factor=self.scale ) - self._unfold_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + 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_stage(self, x, scale, percieve_pattern, stage): b,c,h,w = x.shape @@ -66,10 +107,10 @@ class SRNetR90(nn.Module): 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) - output += self.forward_stage(x, self.scale, self._unfold_pattern_S, self.stage1_S) + output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S) for rotations_count in range(1,4): rotated = torch.rot90(x, k=rotations_count, dims=[2, 3]) - output += torch.rot90(self.forward_stage(x, self.scale, self._unfold_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3]) + output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3]) output /= 4 output = output.reshape(b, c, h*self.scale, w*self.scale) return output @@ -89,7 +130,7 @@ class SRNetR90Y(nn.Module): layers_count=layers_count, upscale_factor=self.scale ) - self._unfold_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=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.rgb_to_ycbcr = layers.RgbToYcbcr() self.ycbcr_to_rgb = layers.YcbcrToRgb() @@ -112,10 +153,10 @@ class SRNetR90Y(nn.Module): x = y.view(b, 1, h, w) output = torch.zeros([b, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device) - output += self.forward_stage(x, self.scale, self._unfold_pattern_S, self.stage1_S) + output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S) 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.scale, self._unfold_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3]) + output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3]) output /= 4 output = torch.cat([output, cbcr_scaled], dim=1) output = self.ycbcr_to_rgb(output).clamp(0, 255)