diff --git a/src/models/__init__.py b/src/models/__init__.py index 9001ba5..6ffe515 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -18,6 +18,10 @@ AVAILABLE_MODELS = { 'SDYLutx1': sdylut.SDYLutx1, 'SDYNetx2': sdynet.SDYNetx2, 'SDYLutx2': sdylut.SDYLutx2, + 'SDYNetR90x1': sdynet.SDYNetR90x1, + 'SDYLutR90x1': sdylut.SDYLutR90x1, + 'SDYNetR90x2': sdynet.SDYNetR90x2, + 'SDYLutR90x2': sdylut.SDYLutR90x2, 'SRNetY': srnet.SRNetY, 'SRLutY': srlut.SRLutY, # 'RCNetCentered_3x3': rcnet.RCNetCentered_3x3, 'RCLutCentered_3x3': rclut.RCLutCentered_3x3, diff --git a/src/models/sdylut.py b/src/models/sdylut.py index f229078..0bb609b 100644 --- a/src/models/sdylut.py +++ b/src/models/sdylut.py @@ -45,6 +45,131 @@ class SDYLutx1(nn.Module): x = x.reshape(b, c, h*scale, w*scale) return x + def forward(self, x): + b,c,h,w = x.shape + x = x.view(b*c, 1, h, w).type(torch.float32) + 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.stageS) + output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stageD) + output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stageY) + output /= 3 + output = round_func(output) + output = output.view(b, c, h*self.scale, w*self.scale) + return output + + def __repr__(self): + return f"{self.__class__.__name__}" + \ + f"\n stageS size: {self.stageS.shape}" + \ + f"\n stageD size: {self.stageD.shape}" + \ + f"\n stageY size: {self.stageY.shape}" + +class SDYLutx2(nn.Module): + def __init__( + self, + quantization_interval, + scale + ): + super(SDYLutx2, self).__init__() + self.scale = scale + self.quantization_interval = quantization_interval + self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3) + self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3) + self.stage1_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stage1_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stage1_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stage2_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stage2_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stage2_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + + @staticmethod + def init_from_numpy( + stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y + ): + scale = int(stageS.shape[-1]) + quantization_interval = 256//(stageS.shape[0]-1) + lut_model = SDYLutx2(quantization_interval=quantization_interval, scale=scale) + lut_model.stage1_S = nn.Parameter(torch.tensor(stage1_S).type(torch.float32)) + lut_model.stage1_D = nn.Parameter(torch.tensor(stage1_D).type(torch.float32)) + lut_model.stage1_Y = nn.Parameter(torch.tensor(stage1_Y).type(torch.float32)) + lut_model.stage2_S = nn.Parameter(torch.tensor(stage2_S).type(torch.float32)) + lut_model.stage2_D = nn.Parameter(torch.tensor(stage2_D).type(torch.float32)) + lut_model.stage2_Y = nn.Parameter(torch.tensor(stage2_Y).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 = x.view(b*c, 1, h, w).type(torch.float32) + 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 /= 3 + output = round_func(output) + 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 /= 3 + output = round_func(output) + output = output.view(b, c, h*self.scale, w*self.scale) + return output + + def __repr__(self): + return f"{self.__class__.__name__}" + \ + f"\n stageS size: {self.stageS.shape}" + \ + f"\n stageD size: {self.stageD.shape}" + \ + f"\n stageY size: {self.stageY.shape}" + +class SDYLutR90x1(nn.Module): + def __init__( + self, + quantization_interval, + scale + ): + super(SDYLutR90x1, self).__init__() + self.scale = scale + self.quantization_interval = quantization_interval + self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) + self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3) + self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3) + self.stageS = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stageD = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + self.stageY = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) + + @staticmethod + def init_from_numpy( + stageS, stageD, stageY + ): + scale = int(stageS.shape[-1]) + quantization_interval = 256//(stageS.shape[0]-1) + lut_model = SDYLutR90x1(quantization_interval=quantization_interval, scale=scale) + lut_model.stageS = nn.Parameter(torch.tensor(stageS).type(torch.float32)) + lut_model.stageD = nn.Parameter(torch.tensor(stageD).type(torch.float32)) + lut_model.stageY = nn.Parameter(torch.tensor(stageY).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 = x.view(b*c, 1, h, w).type(torch.float32) @@ -70,13 +195,13 @@ class SDYLutx1(nn.Module): -class SDYLutx2(nn.Module): +class SDYLutR90x2(nn.Module): def __init__( self, quantization_interval, scale ): - super(SDYLutx2, self).__init__() + super(SDYLutR90x2, self).__init__() self.scale = scale self.quantization_interval = quantization_interval self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2) @@ -95,7 +220,7 @@ class SDYLutx2(nn.Module): ): scale = int(stage2_S.shape[-1]) quantization_interval = 256//(stage2_S.shape[0]-1) - lut_model = SDYLutx2(quantization_interval=quantization_interval, scale=scale) + lut_model = SDYLutR90x2(quantization_interval=quantization_interval, scale=scale) lut_model.stage1_S = nn.Parameter(torch.tensor(stage1_S).type(torch.float32)) lut_model.stage1_D = nn.Parameter(torch.tensor(stage1_D).type(torch.float32)) lut_model.stage1_Y = nn.Parameter(torch.tensor(stage1_Y).type(torch.float32)) diff --git a/src/models/sdynet.py b/src/models/sdynet.py index cc65f37..6791e82 100644 --- a/src/models/sdynet.py +++ b/src/models/sdynet.py @@ -8,6 +8,7 @@ from common import layers from pathlib import Path from . import sdylut + class SDYNetx1(nn.Module): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): super(SDYNetx1, self).__init__() @@ -29,6 +30,101 @@ class SDYNetx1(nn.Module): x = x.reshape(b, c, h*scale, w*scale) return x + def forward(self, x): + 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._extract_pattern_S, self.stage1_S) + output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage1_D) + output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage1_Y) + output /= 3 + x = output + x = round_func(x) + 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): + stageS = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size) + stageD = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size) + stageY = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size) + lut_model = sdylut.SDYLutx1.init_from_numpy(stageS, stageD, stageY) + return lut_model + +class SDYNetx2(nn.Module): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SDYNetx2, 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=1) + self.stage1_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1) + self.stage1_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1) + self.stage2_S = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale) + + 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 = 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 /= 3 + x = output + x = round_func(x) + 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 /= 3 + x = output + x = round_func(x) + 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): + 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.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_D = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size) + stage2_Y = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size) + lut_model = sdylut.SDYLutx2.init_from_numpy(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y) + return lut_model + +class SDYNetR90x1(nn.Module): + def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): + super(SDYNetR90x1, 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) + + 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 = x.reshape(b*c, 1, h, w) @@ -51,12 +147,12 @@ class SDYNetx1(nn.Module): stageS = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size) stageD = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size) stageY = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size) - lut_model = sdylut.SDYLutx1.init_from_numpy(stageS, stageD, stageY) + lut_model = sdylut.SDYLutR90x1.init_from_numpy(stageS, stageD, stageY) return lut_model -class SDYNetx2(nn.Module): +class SDYNetR90x2(nn.Module): def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4): - super(SDYNetx2, self).__init__() + super(SDYNetR90x2, 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) @@ -113,5 +209,5 @@ class SDYNetx2(nn.Module): 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.SDYLutx2.init_from_numpy(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y) + lut_model = sdylut.SDYLutR90x2.init_from_numpy(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y) return lut_model \ No newline at end of file