Spaces:
Running
on
Zero
Running
on
Zero
| from torch import nn, Tensor | |
| from torch.hub import load_state_dict_from_url | |
| from typing import Optional | |
| from .vgg import VGG | |
| from .utils import make_vgg_layers, vgg_urls | |
| from ..utils import _init_weights, ConvDownsample, _get_activation, _get_norm_layer | |
| EPS = 1e-6 | |
| encoder_cfg = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512] | |
| decoder_cfg = [512, 512, 512, 256, 128] | |
| class CSRNet(nn.Module): | |
| def __init__( | |
| self, | |
| model_name: str, | |
| block_size: Optional[int] = None, | |
| norm: str = "none", | |
| act: str = "none" | |
| ) -> None: | |
| super().__init__() | |
| assert model_name in ["vgg16", "vgg16_bn"], f"Model name should be one of ['vgg16', 'vgg16_bn'], but got {model_name}." | |
| assert block_size is None or block_size in [8, 16, 32], f"block_size should be one of [8, 16, 32], but got {block_size}." | |
| self.model_name = model_name | |
| vgg = VGG(make_vgg_layers(encoder_cfg, in_channels=3, batch_norm="bn" in model_name, dilation=1)) | |
| vgg.load_state_dict(load_state_dict_from_url(vgg_urls[model_name]), strict=False) | |
| self.encoder = vgg.features | |
| self.encoder_reduction = 8 | |
| self.encoder_channels = 512 | |
| self.block_size = block_size if block_size is not None else 8 | |
| if norm == "bn": | |
| norm_layer = nn.BatchNorm2d | |
| elif norm == "ln": | |
| norm_layer = nn.LayerNorm | |
| else: | |
| norm_layer = _get_norm_layer(vgg) | |
| if act == "relu": | |
| activation = nn.ReLU(inplace=True) | |
| elif act == "gelu": | |
| activation = nn.GELU() | |
| else: | |
| activation = _get_activation(vgg) | |
| if self.block_size == self.encoder_reduction: | |
| self.refiner = nn.Identity() | |
| elif self.block_size > self.encoder_reduction: | |
| if self.block_size == 32: | |
| self.refiner = nn.Sequential( | |
| ConvDownsample( | |
| in_channels=self.encoder_channels, | |
| out_channels=self.encoder_channels, | |
| norm_layer=norm_layer, | |
| activation=activation, | |
| ), | |
| ConvDownsample( | |
| in_channels=self.encoder_channels, | |
| out_channels=self.encoder_channels, | |
| norm_layer=norm_layer, | |
| activation=activation, | |
| ) | |
| ) | |
| elif self.block_size == 16: | |
| self.refiner = ConvDownsample( | |
| in_channels=self.encoder_channels, | |
| out_channels=self.encoder_channels, | |
| norm_layer=norm_layer, | |
| activation=activation, | |
| ) | |
| self.refiner_channels = self.encoder_channels | |
| self.refiner_reduction = self.block_size | |
| decoder = make_vgg_layers(decoder_cfg, in_channels=512, batch_norm="bn" in model_name, dilation=2) | |
| decoder.apply(_init_weights) | |
| self.decoder = decoder | |
| self.decoder_channels = decoder_cfg[-1] | |
| self.decoder_reduction = self.refiner_reduction | |
| def encode(self, x: Tensor) -> Tensor: | |
| return self.encoder(x) | |
| def refine(self, x: Tensor) -> Tensor: | |
| return self.refiner(x) | |
| def decode(self, x: Tensor) -> Tensor: | |
| return self.decoder(x) | |
| def forward(self, x: Tensor) -> Tensor: | |
| x = self.encode(x) | |
| x = self.refine(x) | |
| x = self.decode(x) | |
| return x | |
| def _csrnet(block_size: Optional[int] = None, norm: str = "none", act: str = "none") -> CSRNet: | |
| return CSRNet("vgg16", block_size=block_size, norm=norm, act=act) | |
| def _csrnet_bn(block_size: Optional[int] = None, norm: str = "none", act: str = "none") -> CSRNet: | |
| return CSRNet("vgg16_bn", block_size=block_size, norm=norm, act=act) | |