左右滑动查看更多
模型搭建介绍
基础模块
def conv(in_channels, out_channels, kernel_size, bias_attr=False, stride=1):
layer = nn.Conv2D(in_channels, out_channels, kernel_size, padding=(kernel_size // 2), bias_attr=bias_attr, stride=stride)
return layer
注意力模块
构建多种注意力模块。
## Spatial Attention
class SALayer(nn.Layer):
def __init__(self, kernel_size=7):
super(SALayer, self).__init__()
self.conv1 = nn.Conv2D(2, 1, kernel_size, padding=kernel_size // 2, bias_attr=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = paddle.mean(x, axis=1, keepdim=True)
max_out = paddle.max(x, axis=1, keepdim=True)
y = paddle.concat([avg_out, max_out], axis=1)
y = self.conv1(y)
y = self.sigmoid(y)
return x * y
# Spatial Attention Block (SAB)
class SAB(nn.Layer):
def __init__(self, n_feat, kernel_size, reduction, bias_attr, act):
super(SAB, self).__init__()
modules_body = [conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr), act, conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr)]
self.body = nn.Sequential(*modules_body)
self.SA = SALayer(kernel_size=7)
def forward(self, x):
res = self.body(x)
res = self.SA(res)
res += x
return res
## Pixel Attention
class PALayer(nn.Layer):
def __init__(self, channel, reduction=16, bias_attr=False):
super(PALayer, self).__init__()
self.pa = nn.Sequential(
nn.Conv2D(channel, channel // reduction, 1, padding=0, bias_attr=bias_attr),
nn.ReLU(),
nn.Conv2D(channel // reduction, channel, 1, padding=0, bias_attr=bias_attr), # channel <-> 1
nn.Sigmoid()
)
def forward(self, x):
y = self.pa(x)
return x * y
## Pixel Attention Block (PAB)
class PAB(nn.Layer):
def __init__(self, n_feat, kernel_size, reduction, bias_attr, act):
super(PAB, self).__init__()
modules_body = [conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr), act, conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr)]
self.PA = PALayer(n_feat, reduction, bias_attr=bias_attr)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.PA(res)
res += x
return res
## Channel Attention Layer
class CALayer(nn.Layer):
def __init__(self, channel, reduction=16, bias_attr=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2D(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2D(channel, channel // reduction, 1, padding=0, bias_attr=bias_attr),
nn.ReLU(),
nn.Conv2D(channel // reduction, channel, 1, padding=0, bias_attr=bias_attr),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x)
y = self.conv_du(y)
return x * y
## Channel Attention Block (CAB)
class CAB(nn.Layer):
def __init__(self, n_feat, kernel_size, reduction, bias_attr, act):
super(CAB, self).__init__()
modules_body = [conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr), act, conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr)]
self.CA = CALayer(n_feat, reduction, bias_attr=bias_attr)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
图像缩放
DownSample：下采样，用于缩小特征图尺寸，提取图像特征；
##---------- Resizing Modules ----------
class DownSample(nn.Layer):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2D(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias_attr=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Layer):
def __init__(self, in_channels, s_factor):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2D(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias_attr=False))
def forward(self, x):
x = self.up(x)
return x
class SkipUpSample(nn.Layer):
def __init__(self, in_channels, s_factor):
super(SkipUpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2D(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias_attr=False))
def forward(self, x, y):
x = self.up(x)
x = x + y
return x
U-Net
使用对称的Encoder和Decoder，对应层级之间相互连接。
SAM模块
SAM（去除了原版图中的Loss，并且调整了其中卷积的核大小）。
# Supervised Attention Module
class SAM(nn.Layer):
def __init__(self, n_feat, kernel_size, bias_attr):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias_attr=bias_attr)
self.conv2 = conv(n_feat, 3, kernel_size, bias_attr=bias_attr)
self.conv3 = conv(3, n_feat, kernel_size, bias_attr=bias_attr)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x2 = nn.functional.sigmoid(self.conv3(img))
x1 = x1 * x2
x1 = x1 + x
return x1, img
MSC模块
Mixed Residual Module实现代码如下：
Mixed Residual Module实现代码如下。
# Mixed Residual Module
class Mix(nn.Layer):
def __init__(self, m=1):
super(Mix, self).__init__()
self.w = self.create_parameter((1,), default_initializer=nn.initializer.Constant(m))
self.mix_block = nn.Sigmoid()
def forward(self, fea1, fea2, feat3):
factor = self.mix_block(self.w)
other = (1 - factor)/2
output = fea1 * other + fea2 * factor + feat3 * other
return output, factor
CMFNet模型
上述的多个模块拼接一下即可搭建出完整的CMFNet：
# CMFNet
class CMFNet(nn.Layer):
def __init__(self, in_c=3, out_c=3, n_feat=96, scale_unetfeats=48, kernel_size=3, reduction=4, bias_attr=False):
super(CMFNet, self).__init__()
p_act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias_attr=bias_attr), p_act,
conv(n_feat // 2, n_feat, kernel_size, bias_attr=bias_attr))
self.shallow_feat2 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias_attr=bias_attr), p_act,
conv(n_feat // 2, n_feat, kernel_size, bias_attr=bias_attr))
self.shallow_feat3 = nn.Sequential(conv(in_c, n_feat // 2, kernel_size, bias_attr=bias_attr), p_act,
conv(n_feat // 2, n_feat, kernel_size, bias_attr=bias_attr))
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'CAB')
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'CAB')
self.stage2_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'PAB')
self.stage2_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'PAB')
self.stage3_encoder = Encoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'SAB')
self.stage3_decoder = Decoder(n_feat, kernel_size, reduction, p_act, bias_attr, scale_unetfeats, 'SAB')
self.sam1o = SAM(n_feat, kernel_size=3, bias_attr=bias_attr)
self.sam2o = SAM(n_feat, kernel_size=3, bias_attr=bias_attr)
self.sam3o = SAM(n_feat, kernel_size=3, bias_attr=bias_attr)
self.mix = Mix(1)
self.add123 = conv(out_c, out_c, kernel_size, bias_attr=bias_attr)
self.concat123 = conv(n_feat*3, n_feat, kernel_size, bias_attr=bias_attr)
self.tail = conv(n_feat, out_c, kernel_size, bias_attr=bias_attr)
def forward(self, x):
## Compute Shallow Features
shallow1 = self.shallow_feat1(x)
shallow2 = self.shallow_feat2(x)
shallow3 = self.shallow_feat3(x)
## Enter the UNet-CAB
x1 = self.stage1_encoder(shallow1)
x1_D = self.stage1_decoder(x1)
## Apply SAM
x1_out, x1_img = self.sam1o(x1_D[0], x)
## Enter the UNet-PAB
x2 = self.stage2_encoder(shallow2)
x2_D = self.stage2_decoder(x2)
## Apply SAM
x2_out, x2_img = self.sam2o(x2_D[0], x)
## Enter the UNet-SAB
x3 = self.stage3_encoder(shallow3)
x3_D = self.stage3_decoder(x3)
## Apply SAM
x3_out, x3_img = self.sam3o(x3_D[0], x)
## Aggregate SAM features of Stage 1, Stage 2 and Stage 3
mix_r = self.mix(x1_img, x2_img, x3_img)
mixed_img = self.add123(mix_r[0])
## Concat SAM features of Stage 1, Stage 2 and Stage 3
concat_feat = self.concat123(paddle.concat([x1_out, x2_out, x3_out], 1))
x_final = self.tail(concat_feat)
return x_final + mixed_img
模型推理过程介绍
功能函数
加载模型：加载训练完成的模型参数；
import cv2
from IPython.display import Image, display
def load_model(model_path):
model = CMFNet()
model.eval()
params = paddle.load(model_path)
model.set_state_dict(params)
return model
def preprocess(img):
clip_h, clip_w = [_ % 4 if _ % 4 else None for _ in img.shape[:2]]
x = img[None, :clip_h, :clip_w, ::-1]
x = x.transpose(0, 3, 1, 2)
x = x.astype('float32')
x /= 255.0
x = paddle.to_tensor(x)
return x
def postprocess(y):
y = y.numpy()
y = y.clip(0.0, 1.0)
y *= 255.0
y = y.transpose(0, 2, 3, 1)
y = y.astype('uint8')
y = y[0, :, :, ::-1]
return y
@paddle.no_grad()
def run(model, img_path, save_path):
img = cv2.imread(img_path)
x = preprocess(img)
y = model(x)
deimg = postprocess(y)
cv2.imwrite(save_path, deimg)
return deimg
def show(img_path, save_path):
display(Image(img_path))
display(Image(save_path))
去模糊
加载去模糊、去雾、去雨水模型，使用上述代码完成模型推理，分别实现去模糊、去雾、去雨水效果。
总结
参考资料
[1] Fan C M, Liu T J, Liu K H. Compound Multi-branch Feature Fusion for Real Image Restoration[J]. arXiv preprint arXiv:2206.02748, 2022.
关注【飞桨PaddlePaddle】公众号
获取更多技术内容~