Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn, Tensor | |
| import torch.nn.functional as F | |
| from typing import List, Optional | |
| from ..utils import _init_weights | |
| from .csrnet import CSRNet, csrnet, csrnet_bn | |
| EPS = 1e-6 | |
| class ContextualModule(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int = 512, | |
| sizes: List[int] = [1, 2, 3, 6], | |
| ) -> None: | |
| super().__init__() | |
| self.scales = nn.ModuleList([self.__make_scale__(in_channels, size) for size in sizes]) | |
| self.bottleneck = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.weight_net = nn.Conv2d(in_channels, in_channels, kernel_size=1) | |
| def __make_weight__(self, feature: Tensor, scale_feature: Tensor) -> Tensor: | |
| weight_feature = feature - scale_feature | |
| weight_feature = self.weight_net(weight_feature) | |
| return F.sigmoid(weight_feature) | |
| def __make_scale__(self, channels: int, size: int) -> nn.Module: | |
| return nn.Sequential( | |
| nn.AdaptiveAvgPool2d(output_size=(size, size)), | |
| nn.Conv2d(channels, channels, kernel_size=1, bias=False), | |
| ) | |
| def forward(self, feature: Tensor) -> Tensor: | |
| h, w = feature.shape[-2:] | |
| multi_scales = [F.interpolate(input=scale(feature), size=(h, w), mode="bilinear") for scale in self.scales] | |
| weights = [self.__make_weight__(feature, scale_feature) for scale_feature in multi_scales] | |
| multi_scales = sum([multi_scales[i] * weights[i] for i in range(len(weights))]) / (sum(weights) + EPS) | |
| overall_features = torch.cat([multi_scales, feature], dim=1) | |
| overall_features = self.bottleneck(overall_features) | |
| overall_features = self.relu(overall_features) | |
| return overall_features | |
| class CANNet(nn.Module): | |
| def __init__( | |
| self, | |
| csrnet: CSRNet, | |
| sizes: List[int] = [1, 2, 3, 6], | |
| reduction: Optional[int] = 8, | |
| ) -> None: | |
| super().__init__() | |
| assert isinstance(csrnet, CSRNet), f"csrnet should be an instance of CSRNet, got {type(csrnet)}." | |
| assert isinstance(sizes, (tuple, list)), f"sizes should be a list or tuple, got {type(sizes)}." | |
| assert len(sizes) > 0, f"Expected at least one size, got {len(sizes)}." | |
| assert all([isinstance(size, int) for size in sizes]), f"Expected all size to be int, got {sizes}." | |
| self.sizes = sizes | |
| self.encoder_reduction = csrnet.encoder_reduction | |
| self.reduction = self.encoder_reduction if reduction is None else reduction | |
| self.features = csrnet.features | |
| self.decoder = csrnet.decoder | |
| self.decoder.apply(_init_weights) | |
| self.context = ContextualModule(512, 512, self.sizes) | |
| self.context.apply(_init_weights) | |
| self.channels = csrnet.channels | |
| def forward(self, x: Tensor) -> Tensor: | |
| x = self.features(x) | |
| x = self.context(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 cannet(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet: | |
| return CANNet(csrnet(), sizes=sizes, reduction=reduction) | |
| def cannet_bn(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet: | |
| return CANNet(csrnet_bn(), sizes=sizes, reduction=reduction) | |