Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- coding:utf-8 -*- | |
| # Power by Zongsheng Yue 2022-07-12 20:35:28 | |
| import math | |
| from torch import nn | |
| import torch.nn.functional as F | |
| class SRCNN(nn.Module): | |
| def __init__(self, in_chns, out_chns=None, num_chns=64, depth=8, sf=4): | |
| super().__init__() | |
| self.sf = sf | |
| out_chns = in_chns if out_chns is None else out_chns | |
| self.head = nn.Conv2d(in_chns, num_chns, kernel_size=5, padding=2) | |
| body = [] | |
| for _ in range(depth-1): | |
| body.append(nn.Conv2d(num_chns, num_chns, kernel_size=5, padding=2)) | |
| body.append(nn.LeakyReLU(0.2, inplace=True)) | |
| self.body = nn.Sequential(*body) | |
| tail = [] | |
| for _ in range(int(math.log(sf, 2))): | |
| tail.append(nn.Conv2d(num_chns, num_chns*4, kernel_size=3, padding=1)) | |
| tail.append(nn.LeakyReLU(0.2, inplace=True)) | |
| tail.append(nn.PixelShuffle(2)) | |
| tail.append(nn.Conv2d(num_chns, out_chns, kernel_size=5, padding=2)) | |
| self.tail = nn.Sequential(*tail) | |
| def forward(self, x): | |
| y = self.head(x) | |
| y = self.body(y) | |
| y = self.tail(y) | |
| return y | |
| class SRCNNFSR(nn.Module): | |
| def __init__(self, in_chns, down_scale_factor=2, num_chns=64, depth=8, sf=4): | |
| super().__init__() | |
| self.sf = sf | |
| head = [] | |
| in_chns_shuffle = in_chns * 4 | |
| assert num_chns % 4 == 0 | |
| for ii in range(int(math.log(down_scale_factor, 2))): | |
| head.append(nn.PixelUnshuffle(2)) | |
| head.append(nn.Conv2d(in_chns_shuffle, num_chns, kernel_size=3, padding=1)) | |
| if ii + 1 < int(math.log(down_scale_factor, 2)): | |
| head.append(nn.Conv2d(num_chns, num_chns//4, kernel_size=5, padding=2)) | |
| head.append(nn.LeakyReLU(0.2, inplace=True)) | |
| in_chns_shuffle = num_chns | |
| self.head = nn.Sequential(*head) | |
| body = [] | |
| for _ in range(depth-1): | |
| body.append(nn.Conv2d(num_chns, num_chns, kernel_size=5, padding=2)) | |
| body.append(nn.LeakyReLU(0.2, inplace=True)) | |
| self.body = nn.Sequential(*body) | |
| tail = [] | |
| for _ in range(int(math.log(down_scale_factor, 2))): | |
| tail.append(nn.Conv2d(num_chns, num_chns, kernel_size=3, padding=1)) | |
| tail.append(nn.LeakyReLU(0.2, inplace=True)) | |
| tail.append(nn.PixelShuffle(2)) | |
| num_chns //= 4 | |
| tail.append(nn.Conv2d(num_chns, in_chns, kernel_size=5, padding=2)) | |
| self.tail = nn.Sequential(*tail) | |
| def forward(self, x): | |
| y = self.head(x) | |
| y = self.body(y) | |
| y = self.tail(y) | |
| return y | |