|
|
@ -5,25 +5,26 @@ import numpy as np
|
|
|
|
from common.utils import round_func
|
|
|
|
from common.utils import round_func
|
|
|
|
from common import lut
|
|
|
|
from common import lut
|
|
|
|
from pathlib import Path
|
|
|
|
from pathlib import Path
|
|
|
|
# from . import srlut
|
|
|
|
from . import hdblut
|
|
|
|
from common import layers
|
|
|
|
from common import layers
|
|
|
|
|
|
|
|
from itertools import cycle
|
|
|
|
|
|
|
|
|
|
|
|
class HDBNet(nn.Module):
|
|
|
|
class HDBNet(nn.Module):
|
|
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
|
|
super(HDBNet, self).__init__()
|
|
|
|
super(HDBNet, self).__init__()
|
|
|
|
assert scale == 4
|
|
|
|
assert scale == 4
|
|
|
|
self.scale = scale
|
|
|
|
self.scale = scale
|
|
|
|
self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2*2)
|
|
|
|
self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2*2)
|
|
|
|
self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2*2)
|
|
|
|
self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2*2)
|
|
|
|
self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2*2)
|
|
|
|
|
|
|
|
|
|
|
|
self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
# self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
# self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
# self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
# self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
# self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2)
|
|
|
|
|
|
|
|
|
|
|
|
self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
|
|
|
|
self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
|
|
|
|
self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
|
|
|
|
self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
|
|
|
@ -41,59 +42,213 @@ class HDBNet(nn.Module):
|
|
|
|
x = x.reshape(b, c, h*scale, w*scale)
|
|
|
|
x = x.reshape(b, c, h*scale, w*scale)
|
|
|
|
return x
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
def forward(self, x, config=None):
|
|
|
|
b,c,h,w = x.shape
|
|
|
|
b,c,h,w = x.shape
|
|
|
|
x = x.reshape(b*c, 1, h, w)
|
|
|
|
x = x.reshape(b*c, 1, h, w)
|
|
|
|
lsb = x % 16
|
|
|
|
lsb = x % 16
|
|
|
|
msb = x - lsb
|
|
|
|
msb = x - lsb
|
|
|
|
output_msb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
|
|
|
|
output_msb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
|
|
|
|
output_lsb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
|
|
|
|
output_lsb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
|
|
|
|
for rotations_count in range(4):
|
|
|
|
for rotations_count in range(4):
|
|
|
|
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2*2, self._extract_pattern_3H, self.stage1_3H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2*2, self._extract_pattern_3D, self.stage1_3D), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2*2, self._extract_pattern_3B, self.stage1_3B), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2*2, self._extract_pattern_2H, self.stage1_2H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2*2, self._extract_pattern_2D, self.stage1_2D), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb /= 4*3
|
|
|
|
output_msb /= 4*3
|
|
|
|
output_lsb /= 4*2
|
|
|
|
output_lsb /= 4*2
|
|
|
|
output_msb = output_msb + output_lsb
|
|
|
|
output_msb = round_func((output_msb / 255) * 16) * 15
|
|
|
|
x = output_msb
|
|
|
|
output_lsb = (output_lsb / 255) * 15
|
|
|
|
lsb = x % 16
|
|
|
|
# print(output_msb.min(), output_msb.max(), output_lsb.min(), output_lsb.max())
|
|
|
|
msb = x - lsb
|
|
|
|
x = output_msb + output_lsb
|
|
|
|
output_msb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
|
|
|
|
# lsb = x % 16
|
|
|
|
output_lsb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
|
|
|
|
# msb = x - lsb
|
|
|
|
for rotations_count in range(4):
|
|
|
|
|
|
|
|
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
# output_msb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
|
|
|
|
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
# output_lsb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H), k=-rotations_count, dims=[2, 3])
|
|
|
|
# for rotations_count in range(4):
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D), k=-rotations_count, dims=[2, 3])
|
|
|
|
# rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B), k=-rotations_count, dims=[2, 3])
|
|
|
|
# rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H), k=-rotations_count, dims=[2, 3])
|
|
|
|
# output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D), k=-rotations_count, dims=[2, 3])
|
|
|
|
# output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb /= 4*3
|
|
|
|
# output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_lsb /= 4*2
|
|
|
|
# output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H), k=-rotations_count, dims=[2, 3])
|
|
|
|
output_msb = output_msb + output_lsb
|
|
|
|
# output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D), k=-rotations_count, dims=[2, 3])
|
|
|
|
x = output_msb
|
|
|
|
# output_msb /= 4*3
|
|
|
|
|
|
|
|
# output_lsb /= 4*2
|
|
|
|
|
|
|
|
# output_msb = round_func((output_msb / 255) * 16) * 15
|
|
|
|
|
|
|
|
# output_lsb = (output_lsb / 255) * 15
|
|
|
|
|
|
|
|
# # print(output_msb.min(), output_msb.max(), output_lsb.min(), output_lsb.max())
|
|
|
|
|
|
|
|
# x = output_msb + output_lsb
|
|
|
|
x = x.reshape(b, c, h*self.scale, w*self.scale)
|
|
|
|
x = x.reshape(b, c, h*self.scale, w*self.scale)
|
|
|
|
return x
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
def get_lut_model(self, quantization_interval=16, batch_size=2**10):
|
|
|
|
def get_lut_model(self, quantization_interval=16, batch_size=2**10):
|
|
|
|
stage1_3H = lut.transfer_2x2_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_3D = lut.transfer_2x2_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_3B = lut.transfer_2x2_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_2H = lut.transfer_2x2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_2D = lut.transfer_2x2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
|
|
|
|
stage2_3H = lut.transfer_2x2_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_3D = lut.transfer_2x2_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_3B = lut.transfer_2x2_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_2H = lut.transfer_2x2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_2D = lut.transfer_2x2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
|
|
|
|
lut_model = hdblut.HDBLut.init_from_numpy(
|
|
|
|
lut_model = hdblut.HDBLut.init_from_numpy(
|
|
|
|
stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D,
|
|
|
|
stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D,
|
|
|
|
stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
|
|
|
|
stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
|
|
|
|
)
|
|
|
|
)
|
|
|
|
return lut_model
|
|
|
|
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 HDBLNet(nn.Module):
|
|
|
|
|
|
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
|
|
|
|
|
|
super(HDBLNet, self).__init__()
|
|
|
|
|
|
|
|
self.scale = scale
|
|
|
|
|
|
|
|
self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
|
|
|
|
|
|
|
|
self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
|
|
|
|
|
|
|
|
self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
|
|
|
|
|
|
|
|
self.stage1_3L = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
|
|
|
|
|
|
|
|
self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
|
|
|
|
|
|
|
|
self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
|
|
|
|
|
|
|
|
self._extract_pattern_3L = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,1]], center=[0,0], window_size=2)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
lsb = x % 16
|
|
|
|
|
|
|
|
msb = x - lsb
|
|
|
|
|
|
|
|
output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
|
|
|
|
|
|
|
|
output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
|
|
|
|
|
|
|
|
for rotations_count in range(4):
|
|
|
|
|
|
|
|
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3H, self.stage1_3H), k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3D, self.stage1_3D), k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb += torch.rot90(self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3B, self.stage1_3B), k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_lsb += torch.rot90(self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_3L, self.stage1_3L), k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb /= 4*3
|
|
|
|
|
|
|
|
output_lsb /= 4
|
|
|
|
|
|
|
|
output_msb = round_func((output_msb / 255) * 16) * 15
|
|
|
|
|
|
|
|
output_lsb = (output_lsb / 255) * 15
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
x = output_msb + output_lsb
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
x = x.reshape(b, c, h*self.scale, w*self.scale)
|
|
|
|
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_loss_fn(self):
|
|
|
|
|
|
|
|
def loss_fn(pred, target):
|
|
|
|
|
|
|
|
return F.mse_loss(pred/255, target/255)
|
|
|
|
|
|
|
|
return loss_fn
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# def get_lut_model(self, quantization_interval=16, batch_size=2**10):
|
|
|
|
|
|
|
|
# stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
# stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# lut_model = hdblut.HDBLut.init_from_numpy(
|
|
|
|
|
|
|
|
# stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D,
|
|
|
|
|
|
|
|
# stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
|
|
|
|
|
|
|
|
# )
|
|
|
|
|
|
|
|
# return lut_model
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class HDBHNet(nn.Module):
|
|
|
|
|
|
|
|
def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
|
|
|
|
|
|
|
|
super(HDBHNet, self).__init__()
|
|
|
|
|
|
|
|
self.scale = scale
|
|
|
|
|
|
|
|
self.hidden_dim = hidden_dim
|
|
|
|
|
|
|
|
self.layers_count = layers_count
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
self.msb_fns = nn.ModuleList([layers.UpscaleBlock(
|
|
|
|
|
|
|
|
in_features=4,
|
|
|
|
|
|
|
|
hidden_dim=hidden_dim,
|
|
|
|
|
|
|
|
layers_count=layers_count,
|
|
|
|
|
|
|
|
upscale_factor=self.scale
|
|
|
|
|
|
|
|
) for x in range(1)])
|
|
|
|
|
|
|
|
self.lsb_fns = nn.ModuleList([layers.UpscaleBlock(
|
|
|
|
|
|
|
|
in_features=4,
|
|
|
|
|
|
|
|
hidden_dim=hidden_dim,
|
|
|
|
|
|
|
|
layers_count=layers_count,
|
|
|
|
|
|
|
|
upscale_factor=self.scale
|
|
|
|
|
|
|
|
) for x in range(1)])
|
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
lsb = x % 16
|
|
|
|
|
|
|
|
msb = x - lsb
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
|
|
|
|
|
|
|
|
output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
|
|
|
|
|
|
|
|
for rotations_count, msb_fn, lsb_fn in zip(range(4), cycle(self.msb_fns), cycle(self.lsb_fns)):
|
|
|
|
|
|
|
|
rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb_r = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, msb_fn)
|
|
|
|
|
|
|
|
output_lsb_r = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, lsb_fn)
|
|
|
|
|
|
|
|
output_msb_r = round_func((output_msb_r / 255)*16) * 15
|
|
|
|
|
|
|
|
output_lsb_r = (output_lsb_r / 255) * 15
|
|
|
|
|
|
|
|
output_msb += torch.rot90(output_msb_r, k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_lsb += torch.rot90(output_lsb_r, k=-rotations_count, dims=[2, 3])
|
|
|
|
|
|
|
|
output_msb /= 4
|
|
|
|
|
|
|
|
output_lsb /= 4
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
x = output_msb + output_lsb
|
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_loss_fn(self):
|
|
|
|
|
|
|
|
fourier_loss_fn = FocalFrequencyLoss()
|
|
|
|
|
|
|
|
high_frequency_loss_fn = FourierLoss()
|
|
|
|
|
|
|
|
def loss_fn(pred, target):
|
|
|
|
|
|
|
|
a = fourier_loss_fn(pred/255, target/255) * 1e8
|
|
|
|
|
|
|
|
# b = F.mse_loss(pred/255, target/255) #* 1e3
|
|
|
|
|
|
|
|
# c = high_frequency_loss_fn(pred/255, target/255) * 1e6
|
|
|
|
|
|
|
|
return a #+ b #+ c
|
|
|
|
|
|
|
|
return loss_fn
|