models/ebc/hrnet.py

import timm
import torch.nn.functional as F
from torch import nn, Tensor
from functools import partial
from typing import Optional

from ..utils import ConvRefine, _get_norm_layer, _get_activation


available_hrnets = [
    "hrnet_w18", "hrnet_w18_small", "hrnet_w18_small_v2",
    "hrnet_w30", "hrnet_w32", "hrnet_w40", "hrnet_w44", "hrnet_w48", "hrnet_w64",
]


class HRNet(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 available_hrnets, f"Model name should be one of {available_hrnets}"
        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
        self.block_size = block_size if block_size is not None else 32

        model = timm.create_model(model_name, pretrained=True)

        self.conv1 = model.conv1
        self.bn1 = model.bn1
        self.act1 = model.act1
        self.conv2 = model.conv2
        self.bn2 = model.bn2
        self.act2 = model.act2

        self.layer1 = model.layer1

        self.transition1 = model.transition1
        self.stage2 = model.stage2

        self.transition2 = model.transition2
        self.stage3 = model.stage3

        self.transition3 = model.transition3
        self.stage4 = model.stage4

        incre_modules = model.incre_modules
        downsamp_modules = model.downsamp_modules

        assert len(incre_modules) == 4, f"Expected 4 incre_modules, got {len(self.incre_modules)}"
        assert len(downsamp_modules) == 3, f"Expected 3 downsamp_modules, got {len(self.downsamp_modules)}"

        self.out_channels_4 = incre_modules[0][0].downsample[0].out_channels
        self.out_channels_8 = incre_modules[1][0].downsample[0].out_channels
        self.out_channels_16 = incre_modules[2][0].downsample[0].out_channels
        self.out_channels_32 = incre_modules[3][0].downsample[0].out_channels

        if self.block_size == 8:
            self.encoder_reduction = 8
            self.encoder_channels = self.out_channels_8
            self.incre_modules = incre_modules[:2]
            self.downsamp_modules = downsamp_modules[:1]

            self.refiner = nn.Identity()
            self.refiner_reduction = 8
            self.refiner_channels = self.out_channels_8
        
        elif self.block_size == 16:
            self.encoder_reduction = 16
            self.encoder_channels = self.out_channels_16
            self.incre_modules = incre_modules[:3]
            self.downsamp_modules = downsamp_modules[:2]

            self.refiner = nn.Identity()
            self.refiner_reduction = 16
            self.refiner_channels = self.out_channels_16

        else:  # self.block_size == 32
            self.encoder_reduction = 32
            self.encoder_channels = self.out_channels_32
            self.incre_modules = incre_modules
            self.downsamp_modules = downsamp_modules

            self.refiner = nn.Identity()
            self.refiner_reduction = 32
            self.refiner_channels = self.out_channels_32

        # define the decoder
        if self.refiner_channels <= 512:
            groups = 1
        elif self.refiner_channels <= 1024:
            groups = 2
        elif self.refiner_channels <= 2048:
            groups = 4
        else:
            groups = 8
        
        if norm == "bn":
            norm_layer = nn.BatchNorm2d
        elif norm == "ln":
            norm_layer = nn.LayerNorm
        else:
            norm_layer = _get_norm_layer(model)

        if act == "relu":
            activation = nn.ReLU(inplace=True)
        elif act == "gelu":
            activation = nn.GELU()
        else:
            activation = _get_activation(model)
        
        decoder_block = partial(ConvRefine, groups=groups, norm_layer=norm_layer, activation=activation)
        if self.refiner_channels <= 256:
            self.decoder = nn.Identity()
            self.decoder_channels = self.refiner_channels
        elif self.refiner_channels <= 512:
            self.decoder = decoder_block(self.refiner_channels, self.refiner_channels // 2)
            self.decoder_channels = self.refiner_channels // 2
        elif self.refiner_channels <= 1024:
            self.decoder = nn.Sequential(
                decoder_block(self.refiner_channels, self.refiner_channels // 2),
                decoder_block(self.refiner_channels // 2, self.refiner_channels // 4),
            )
            self.decoder_channels = self.refiner_channels // 4
        else:
            self.decoder = nn.Sequential(
                decoder_block(self.refiner_channels, self.refiner_channels // 2),
                decoder_block(self.refiner_channels // 2, self.refiner_channels // 4),
                decoder_block(self.refiner_channels // 4, self.refiner_channels // 8),
            )
            self.decoder_channels = self.refiner_channels // 8
        
        self.decoder_reduction = self.refiner_reduction

    def _interpolate(self, x: Tensor) -> Tensor:
        # This method adjust the spatial dimensions of the input tensor so that it can be divided by 32.
        if x.shape[-1] % 32 != 0 or x.shape[-2] % 32 != 0:
            new_h = int(round(x.shape[-2] / 32) * 32)
            new_w = int(round(x.shape[-1] / 32) * 32)
            return F.interpolate(x, size=(new_h, new_w), mode="bicubic", align_corners=False)
        
        return x


    def encode(self, x: Tensor) -> Tensor:
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.act1(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.act2(x)

        x = self.layer1(x)

        x = [t(x) for t in self.transition1]
        x = self.stage2(x)

        x = [t(x[-1]) if not isinstance(t, nn.Identity) else x[i] for i, t in enumerate(self.transition2)]
        x = self.stage3(x)

        x = [t(x[-1]) if not isinstance(t, nn.Identity) else x[i] for i, t in enumerate(self.transition3)]
        x = self.stage4(x)

        assert len(x) == 4, f"Expected 4 outputs, got {len(x)}"

        feats = None
        for i, incre in enumerate(self.incre_modules):
            if feats is None:
                feats = incre(x[i])
            else:
                down = self.downsamp_modules[i - 1]  # needed for torchscript module indexing
                feats = incre(x[i]) + down.forward(feats)
        
        return feats

    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._interpolate(x)
        x = self.encode(x)
        x = self.refine(x)
        x = self.decode(x)
        return x


def _hrnet(model_name: str, block_size: Optional[int] = None, norm: str = "none", act: str = "none") -> HRNet:
    return HRNet(model_name, block_size, norm, act)