bugfix, new model

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,

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__(

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)