voitl/unet_plus_plus

1import segmentation_models_pytorch as smp2from .hf_config import UnetConfig3from transformers import PreTrainedModel4 5 6class HFUnetPlusPlus(PreTrainedModel):7    config_class = UnetConfig8 9    def __init__(self, config):10        super().__init__(config)11 12        self.model = smp.UnetPlusPlus(13            encoder_name=config.encoder_name,14            encoder_weights="imagenet",15            decoder_channels=config.decoder_channels,16            in_channels=config.input_channels,17            classes=config.num_classes,18            decoder_attention_type="scse")19 20    def forward(self, tensor):21        return self.model(tensor)22