Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn as nn | |
| from basicsr.utils.registry import ARCH_REGISTRY | |
| from .arch_util import Upsample, make_layer | |
| class ChannelAttention(nn.Module): | |
| """Channel attention used in RCAN. | |
| Args: | |
| num_feat (int): Channel number of intermediate features. | |
| squeeze_factor (int): Channel squeeze factor. Default: 16. | |
| """ | |
| def __init__(self, num_feat, squeeze_factor=16): | |
| super(ChannelAttention, self).__init__() | |
| self.attention = nn.Sequential( | |
| nn.AdaptiveAvgPool2d(1), nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0), | |
| nn.ReLU(inplace=True), nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0), nn.Sigmoid()) | |
| def forward(self, x): | |
| y = self.attention(x) | |
| return x * y | |
| class RCAB(nn.Module): | |
| """Residual Channel Attention Block (RCAB) used in RCAN. | |
| Args: | |
| num_feat (int): Channel number of intermediate features. | |
| squeeze_factor (int): Channel squeeze factor. Default: 16. | |
| res_scale (float): Scale the residual. Default: 1. | |
| """ | |
| def __init__(self, num_feat, squeeze_factor=16, res_scale=1): | |
| super(RCAB, self).__init__() | |
| self.res_scale = res_scale | |
| self.rcab = nn.Sequential( | |
| nn.Conv2d(num_feat, num_feat, 3, 1, 1), nn.ReLU(True), nn.Conv2d(num_feat, num_feat, 3, 1, 1), | |
| ChannelAttention(num_feat, squeeze_factor)) | |
| def forward(self, x): | |
| res = self.rcab(x) * self.res_scale | |
| return res + x | |
| class ResidualGroup(nn.Module): | |
| """Residual Group of RCAB. | |
| Args: | |
| num_feat (int): Channel number of intermediate features. | |
| num_block (int): Block number in the body network. | |
| squeeze_factor (int): Channel squeeze factor. Default: 16. | |
| res_scale (float): Scale the residual. Default: 1. | |
| """ | |
| def __init__(self, num_feat, num_block, squeeze_factor=16, res_scale=1): | |
| super(ResidualGroup, self).__init__() | |
| self.residual_group = make_layer( | |
| RCAB, num_block, num_feat=num_feat, squeeze_factor=squeeze_factor, res_scale=res_scale) | |
| self.conv = nn.Conv2d(num_feat, num_feat, 3, 1, 1) | |
| def forward(self, x): | |
| res = self.conv(self.residual_group(x)) | |
| return res + x | |
| class RCAN(nn.Module): | |
| """Residual Channel Attention Networks. | |
| ``Paper: Image Super-Resolution Using Very Deep Residual Channel Attention Networks`` | |
| Reference: https://github.com/yulunzhang/RCAN | |
| Args: | |
| num_in_ch (int): Channel number of inputs. | |
| num_out_ch (int): Channel number of outputs. | |
| num_feat (int): Channel number of intermediate features. | |
| Default: 64. | |
| num_group (int): Number of ResidualGroup. Default: 10. | |
| num_block (int): Number of RCAB in ResidualGroup. Default: 16. | |
| squeeze_factor (int): Channel squeeze factor. Default: 16. | |
| upscale (int): Upsampling factor. Support 2^n and 3. | |
| Default: 4. | |
| res_scale (float): Used to scale the residual in residual block. | |
| Default: 1. | |
| img_range (float): Image range. Default: 255. | |
| rgb_mean (tuple[float]): Image mean in RGB orders. | |
| Default: (0.4488, 0.4371, 0.4040), calculated from DIV2K dataset. | |
| """ | |
| def __init__(self, | |
| num_in_ch, | |
| num_out_ch, | |
| num_feat=64, | |
| num_group=10, | |
| num_block=16, | |
| squeeze_factor=16, | |
| upscale=4, | |
| res_scale=1, | |
| img_range=255., | |
| rgb_mean=(0.4488, 0.4371, 0.4040)): | |
| super(RCAN, self).__init__() | |
| self.img_range = img_range | |
| self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) | |
| self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) | |
| self.body = make_layer( | |
| ResidualGroup, | |
| num_group, | |
| num_feat=num_feat, | |
| num_block=num_block, | |
| squeeze_factor=squeeze_factor, | |
| res_scale=res_scale) | |
| self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) | |
| self.upsample = Upsample(upscale, num_feat) | |
| self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) | |
| def forward(self, x): | |
| self.mean = self.mean.type_as(x) | |
| x = (x - self.mean) * self.img_range | |
| x = self.conv_first(x) | |
| res = self.conv_after_body(self.body(x)) | |
| res += x | |
| x = self.conv_last(self.upsample(res)) | |
| x = x / self.img_range + self.mean | |
| return x | |