You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
259 lines
18 KiB
Python
259 lines
18 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from common.utils import round_func
|
|
from common import lut
|
|
from common import layers
|
|
from pathlib import Path
|
|
from . import sdylut
|
|
from models.base import SRNetBase
|
|
|
|
class SDYNetx1(SRNetBase):
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
super(SDYNetx1, 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(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)
|
|
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 = 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
|
|
|
|
def get_loss_fn(self):
|
|
def loss_fn(pred, target):
|
|
return F.mse_loss(pred/255, target/255)
|
|
return loss_fn
|
|
|
|
|
|
class SDYNetx2(SRNetBase):
|
|
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(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, 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
|
|
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 = 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.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)
|
|
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 SDYNetx3(SRNetBase):
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
super(SDYNetx3, 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=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)
|
|
self.stage3_S = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
|
|
self.stage3_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
|
|
self.stage3_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
|
|
|
|
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, 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
|
|
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.stage2_S)
|
|
output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage2_D)
|
|
output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage2_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.stage3_S)
|
|
output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage3_D)
|
|
output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage3_Y)
|
|
output /= 3
|
|
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):
|
|
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)
|
|
stage3_S = lut.transfer_2x2_input_SxS_output(self.stage3_S, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
stage3_D = lut.transfer_2x2_input_SxS_output(self.stage3_D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
stage3_Y = lut.transfer_2x2_input_SxS_output(self.stage3_Y, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
lut_model = sdylut.SDYLutx3.init_from_numpy(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y, stage3_S, stage3_D, stage3_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 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, 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)
|
|
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)
|
|
for rotations_count in range(1, 4):
|
|
rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
|
|
output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[-2, -1])
|
|
output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1])
|
|
output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])
|
|
output /= 4*3
|
|
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):
|
|
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.SDYLutR90x1.init_from_numpy(stageS, stageD, stageY)
|
|
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 SDYNetR90x2(SRNetBase):
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
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)
|
|
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(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, w], dtype=x.dtype, device=x.device)
|
|
output_1 += self.forward_stage(x, 1, self._extract_pattern_S, self.stage1_S)
|
|
output_1 += self.forward_stage(x, 1, self._extract_pattern_D, self.stage1_D)
|
|
output_1 += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
|
|
for rotations_count in range(1,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 /= 4*3
|
|
x = output_1
|
|
output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
|
|
output_2 += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S)
|
|
output_2 += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D)
|
|
output_2 += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y)
|
|
for rotations_count in range(1,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 /= 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 |