|
|
import os
|
|
|
|
|
|
from .utils import MODELS, imagenet_weights
|
|
|
from .utils import tome_presets
|
|
|
from .model.base_module import BaseModule
|
|
|
from .configs.config.config import Config
|
|
|
from .utils.build_functions import build_model_from_cfg
|
|
|
|
|
|
|
|
|
class SegFormer(BaseModule):
|
|
|
"""
|
|
|
This class represents a SegFormer model that allows for the application of token merging.
|
|
|
|
|
|
Attributes:
|
|
|
backbone (BaseModule): MixVisionTransformer backbone
|
|
|
decode_head (BaseModule): SegFormer head
|
|
|
|
|
|
"""
|
|
|
def __init__(self, cfg):
|
|
|
"""
|
|
|
Initialize the SegFormer model.
|
|
|
|
|
|
Args:
|
|
|
cfg (Config): an mmengine Config object, which defines the backbone, head and token merging strategy used.
|
|
|
|
|
|
"""
|
|
|
super().__init__()
|
|
|
self.backbone = build_model_from_cfg(cfg.backbone, registry=MODELS)
|
|
|
self.decode_head = build_model_from_cfg(cfg.decode_head, registry=MODELS)
|
|
|
|
|
|
def forward(self, x):
|
|
|
"""
|
|
|
Forward pass of the model.
|
|
|
|
|
|
Args:
|
|
|
x (torch.Tensor): input tensor of shape [B, C, H, W]
|
|
|
|
|
|
Returns:
|
|
|
torch.Tensor: output tensor
|
|
|
|
|
|
"""
|
|
|
x = self.backbone(x)
|
|
|
x = self.decode_head(x)
|
|
|
return x
|
|
|
|
|
|
|
|
|
def create_model(
|
|
|
backbone: str = 'b0',
|
|
|
tome_strategy: str = None,
|
|
|
out_channels: int = 19,
|
|
|
pretrained: bool = False,
|
|
|
):
|
|
|
"""
|
|
|
Create a SegFormer model using the predefined SegFormer backbones from the MiT series (b0-b5).
|
|
|
|
|
|
Args:
|
|
|
backbone (str): backbone name (e.g. 'b0')
|
|
|
tome_strategy (str | list(dict)): select strategy from presets ('bsm_hq', 'bsm_fast', 'n2d_2x2') or define a
|
|
|
custom strategy using a list, that contains of dictionaries, in which the strategies for the stage are
|
|
|
defined
|
|
|
out_channels (int): number of output channels (e.g. 19 for the cityscapes semantic segmentation task)
|
|
|
pretrained: use pretrained (imagenet) weights
|
|
|
|
|
|
Returns:
|
|
|
BaseModule: SegFormer model
|
|
|
|
|
|
"""
|
|
|
backbone = backbone.lower()
|
|
|
assert backbone in [f'b{i}' for i in range(6)]
|
|
|
|
|
|
wd = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
cfg = Config.fromfile(os.path.join(wd, 'configs', f'segformer_mit_{backbone}.py'))
|
|
|
|
|
|
cfg.decode_head.out_channels = out_channels
|
|
|
|
|
|
if tome_strategy is not None:
|
|
|
if tome_strategy not in list(tome_presets.keys()):
|
|
|
print("Using custom merging strategy.")
|
|
|
cfg.backbone.tome_cfg = tome_presets[tome_strategy]
|
|
|
|
|
|
|
|
|
if pretrained:
|
|
|
cfg.backbone.init_cfg = dict(type='Pretrained', checkpoint=imagenet_weights[backbone])
|
|
|
|
|
|
return SegFormer(cfg)
|
|
|
|
|
|
|
|
|
def create_custom_model(
|
|
|
model_cfg: Config,
|
|
|
tome_strategy: list[dict] = None,
|
|
|
):
|
|
|
"""
|
|
|
Create a SegFormer model with customizable backbone and head.
|
|
|
|
|
|
Args:
|
|
|
model_cfg (Config): backbone name (e.g. 'b0')
|
|
|
tome_strategy (list(dict)): custom token merging strategy
|
|
|
|
|
|
Returns:
|
|
|
BaseModule: SegFormer model
|
|
|
|
|
|
"""
|
|
|
if tome_strategy is not None:
|
|
|
model_cfg.backbone.tome_cfg = tome_strategy
|
|
|
|
|
|
return SegFormer(model_cfg)
|
|
|
|
