franchesoni/segmentation_features

1print("Importing external...")2import torch3from torch import nn4import torch.nn.functional as F5 6from timm.models.efficientvit_mit import (7    ConvNormAct,8    FusedMBConv,9    MBConv,10    ResidualBlock,11    efficientvit_l1,12)13from timm.layers import GELUTanh14 15 16def val2list(x: list or tuple or any, repeat_time=1):17    if isinstance(x, (list, tuple)):18        return list(x)19    return [x for _ in range(repeat_time)]20 21 22def resize(23    x: torch.Tensor,24    size: any or None = None,25    scale_factor: list[float] or None = None,26    mode: str = "bicubic",27    align_corners: bool or None = False,28) -> torch.Tensor:29    if mode in {"bilinear", "bicubic"}:30        return F.interpolate(31            x,32            size=size,33            scale_factor=scale_factor,34            mode=mode,35            align_corners=align_corners,36        )37    elif mode in {"nearest", "area"}:38        return F.interpolate(x, size=size, scale_factor=scale_factor, mode=mode)39    else:40        raise NotImplementedError(f"resize(mode={mode}) not implemented.")41 42 43class UpSampleLayer(nn.Module):44    def __init__(45        self,46        mode="bicubic",47        size: int or tuple[int, int] or list[int] or None = None,48        factor=2,49        align_corners=False,50    ):51        super(UpSampleLayer, self).__init__()52        self.mode = mode53        self.size = val2list(size, 2) if size is not None else None54        self.factor = None if self.size is not None else factor55        self.align_corners = align_corners56 57    def forward(self, x: torch.Tensor) -> torch.Tensor:58        if (59            self.size is not None and tuple(x.shape[-2:]) == self.size60        ) or self.factor == 1:61            return x62        return resize(x, self.size, self.factor, self.mode, self.align_corners)63 64 65class DAGBlock(nn.Module):66    def __init__(67        self,68        inputs: dict[str, nn.Module],69        merge: str,70        post_input: nn.Module or None,71        middle: nn.Module,72        outputs: dict[str, nn.Module],73    ):74        super(DAGBlock, self).__init__()75 76        self.input_keys = list(inputs.keys())77        self.input_ops = nn.ModuleList(list(inputs.values()))78        self.merge = merge79        self.post_input = post_input80 81        self.middle = middle82 83        self.output_keys = list(outputs.keys())84        self.output_ops = nn.ModuleList(list(outputs.values()))85 86    def forward(self, feature_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:87        feat = [88            op(feature_dict[key]) for key, op in zip(self.input_keys, self.input_ops)89        ]90        if self.merge == "add":91            feat = list_sum(feat)92        elif self.merge == "cat":93            feat = torch.concat(feat, dim=1)94        else:95            raise NotImplementedError96        if self.post_input is not None:97            feat = self.post_input(feat)98        feat = self.middle(feat)99        for key, op in zip(self.output_keys, self.output_ops):100            feature_dict[key] = op(feat)101        return feature_dict102 103 104def list_sum(x: list) -> any:105    return x[0] if len(x) == 1 else x[0] + list_sum(x[1:])106 107 108class SegHead(nn.Module):109    def __init__(110        self,111        fid_list: list[str],112        in_channel_list: list[int],113        stride_list: list[int],114        head_stride: int,115        head_width: int,116        head_depth: int,117        expand_ratio: float,118        middle_op: str,119        final_expand: float or None,120        n_classes: int,121        dropout=0,122        norm="bn2d",123        act_func="hswish",124    ):125        super(SegHead, self).__init__()126        # exceptions to adapt effvit to timm127        if act_func == "gelu":128            act_func = GELUTanh129        else:130            raise ValueError(f"act_func {act_func} not supported")131        if norm == "bn2d":132            norm_layer = nn.BatchNorm2d133        else:134            raise ValueError(f"norm {norm} not supported")135 136        inputs = {}137        for fid, in_channel, stride in zip(fid_list, in_channel_list, stride_list):138            factor = stride // head_stride139            if factor == 1:140                inputs[fid] = ConvNormAct(141                    in_channel, head_width, 1, norm_layer=norm_layer, act_layer=act_func142                )143            else:144                inputs[fid] = nn.Sequential(145                    ConvNormAct(146                        in_channel,147                        head_width,148                        1,149                        norm_layer=norm_layer,150                        act_layer=act_func,151                    ),152                    UpSampleLayer(factor=factor),153                )154        self.in_keys = inputs.keys()155        self.in_ops = nn.ModuleList(inputs.values())156 157        middle = []158        for _ in range(head_depth):159            if middle_op == "mbconv":160                block = MBConv(161                    head_width,162                    head_width,163                    expand_ratio=expand_ratio,164                    norm_layer=norm_layer,165                    act_layer=(act_func, act_func, None),166                )167            elif middle_op == "fmbconv":168                block = FusedMBConv(169                    head_width,170                    head_width,171                    expand_ratio=expand_ratio,172                    norm_layer=norm_layer,173                    act_layer=(act_func, None),174                )175            else:176                raise NotImplementedError177            middle.append(ResidualBlock(block, nn.Identity()))178        self.middle = nn.Sequential(*middle)179 180        self.out_layer = nn.Sequential(181            *[182                None183                if final_expand is None184                else ConvNormAct(185                    head_width,186                    head_width * final_expand,187                    1,188                    norm_layer=norm_layer,189                    act_layer=act_func,190                ),191                ConvNormAct(192                    head_width * (final_expand or 1),193                    n_classes,194                    1,195                    bias=True,196                    dropout=dropout,197                    norm_layer=None,198                    act_layer=None,199                ),200            ]201        )202 203    def forward(self, feature_map_list):204        t_feat_maps = [205            self.in_ops[ind](feature_map_list[ind])206            for ind in range(len(feature_map_list))207        ]208        t_feat_map = list_sum(t_feat_maps)209        t_feat_map = self.middle(t_feat_map)210        out = self.out_layer(t_feat_map)211        return out212 213 214class EfficientViT_l1_r224(nn.Module):215    def __init__(216        self,217        out_channels,218        out_ds_factor=1,219        decoder_size="small",220        pretrained=False,221        use_norm_params=False,222    ):223        if decoder_size == "small":224            head_width = 32225            head_depth = 1226            middle_op = "mbconv"227        elif decoder_size == "medium":228            head_width = 64229            head_depth = 3230            middle_op = "mbconv"231        elif decoder_size == "large":232            head_width = 256233            head_depth = 3234            middle_op = "fmbconv"235 236        super(EfficientViT_l1_r224, self).__init__()237        self.bbone = efficientvit_l1(238            num_classes=0, features_only=True, pretrained=pretrained239        )240        self.head = SegHead(241            fid_list=["stage4", "stage3", "stage2"],242            in_channel_list=[512, 256, 128],243            stride_list=[32, 16, 8],244            head_stride=out_ds_factor,245            head_width=head_width,246            head_depth=head_depth,247            expand_ratio=4,248            middle_op=middle_op,249            final_expand=8,250            n_classes=out_channels,251            act_func="gelu",252        )253        # [optional] deactivate normalization254        if not use_norm_params:255            for module in self.modules():256                if (257                    isinstance(module, nn.LayerNorm)258                    or isinstance(module, nn.BatchNorm2d)259                    or isinstance(module, nn.BatchNorm1d)260                ):261                    module.weight.requires_grad_(False)262                    module.bias.requires_grad_(False)263 264    def forward(self, x):265        feat = self.bbone(x)266        out = self.head([feat[3], feat[2], feat[1]])267        return out268