File size: 3,792 Bytes
36d9761 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
from torch import nn as nn
from basicsr.archs.arch_util import ResidualBlockNoBN, default_init_weights
from basicsr.utils.registry import ARCH_REGISTRY
@ARCH_REGISTRY.register()
class DEResNet(nn.Module):
"""Degradation Estimator with ResNetNoBN arch. v2.1, no vector anymore
As shown in paper 'Towards Flexible Blind JPEG Artifacts Removal',
resnet arch works for image quality estimation.
Args:
num_in_ch (int): channel number of inputs. Default: 3.
num_degradation (int): num of degradation the DE should estimate. Default: 2(blur+noise).
degradation_embed_size (int): embedding size of each degradation vector.
degradation_degree_actv (int): activation function for degradation degree scalar. Default: sigmoid.
num_feats (list): channel number of each stage.
num_blocks (list): residual block of each stage.
downscales (list): downscales of each stage.
"""
def __init__(self,
num_in_ch=3,
num_degradation=2,
degradation_degree_actv='sigmoid',
num_feats=(64, 128, 256, 512),
num_blocks=(2, 2, 2, 2),
downscales=(2, 2, 2, 1)):
super(DEResNet, self).__init__()
assert isinstance(num_feats, list)
assert isinstance(num_blocks, list)
assert isinstance(downscales, list)
assert len(num_feats) == len(num_blocks) and len(num_feats) == len(downscales)
num_stage = len(num_feats)
self.conv_first = nn.ModuleList()
for _ in range(num_degradation):
self.conv_first.append(nn.Conv2d(num_in_ch, num_feats[0], 3, 1, 1))
self.body = nn.ModuleList()
for _ in range(num_degradation):
body = list()
for stage in range(num_stage):
for _ in range(num_blocks[stage]):
body.append(ResidualBlockNoBN(num_feats[stage]))
if downscales[stage] == 1:
if stage < num_stage - 1 and num_feats[stage] != num_feats[stage + 1]:
body.append(nn.Conv2d(num_feats[stage], num_feats[stage + 1], 3, 1, 1))
continue
elif downscales[stage] == 2:
body.append(nn.Conv2d(num_feats[stage], num_feats[min(stage + 1, num_stage - 1)], 3, 2, 1))
else:
raise NotImplementedError
self.body.append(nn.Sequential(*body))
# self.body = nn.Sequential(*body)
self.num_degradation = num_degradation
self.fc_degree = nn.ModuleList()
if degradation_degree_actv == 'sigmoid':
actv = nn.Sigmoid
elif degradation_degree_actv == 'tanh':
actv = nn.Tanh
else:
raise NotImplementedError(f'only sigmoid and tanh are supported for degradation_degree_actv, '
f'{degradation_degree_actv} is not supported yet.')
for _ in range(num_degradation):
self.fc_degree.append(
nn.Sequential(
nn.Linear(num_feats[-1], 512),
nn.ReLU(inplace=True),
nn.Linear(512, 1),
actv(),
))
self.avg_pool = nn.AdaptiveAvgPool2d(1)
default_init_weights([self.conv_first, self.body, self.fc_degree], 0.1)
def forward(self, x):
degrees = []
for i in range(self.num_degradation):
x_out = self.conv_first[i](x)
feat = self.body[i](x_out)
feat = self.avg_pool(feat)
feat = feat.squeeze(-1).squeeze(-1)
# for i in range(self.num_degradation):
degrees.append(self.fc_degree[i](feat).squeeze(-1))
return degrees
|