| from torch import nn, Tensor |
| import torch.nn.functional as F |
| from typing import Optional |
|
|
| from ..utils import _init_weights, make_vgg_layers, vgg_urls |
| from .vgg import _load_weights |
|
|
| 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, 64] |
|
|
|
|
| class CSRNet(nn.Module): |
| def __init__( |
| self, |
| features: nn.Module, |
| decoder: nn.Module, |
| reduction: Optional[int] = None, |
| ) -> None: |
| super().__init__() |
| self.features = features |
| self.features.apply(_init_weights) |
| self.decoder = decoder |
| self.decoder.apply(_init_weights) |
|
|
| self.encoder_reduction = 8 |
| self.reduction = self.encoder_reduction if reduction is None else reduction |
| self.channels = 64 |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x = self.features(x) |
| if self.encoder_reduction != self.reduction: |
| x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear") |
| x = self.decoder(x) |
| return x |
|
|
|
|
| def csrnet(reduction: int = 8) -> CSRNet: |
| model = CSRNet( |
| make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=False, dilation=1), |
| make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=False, dilation=2), |
| reduction=reduction |
| ) |
| return _load_weights(model, vgg_urls["vgg16"]) |
|
|
| def csrnet_bn(reduction: int = 8) -> CSRNet: |
| model = CSRNet( |
| make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=True, dilation=1), |
| make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=True, dilation=2), |
| reduction=reduction |
| ) |
| return _load_weights(model, vgg_urls["vgg16"]) |
|
|
