import torch.nn as nn

from isegm.model.modifiers import LRMult
from isegm.utils.serialization import serialize

from .is_model import ISModel
from .modeling.basic_blocks import SepConvHead
from .modeling.deeplab_v3 import DeepLabV3Plus


class DeeplabModel(ISModel):
    @serialize
    def __init__(
        self,
        backbone="resnet50",
        deeplab_ch=256,
        aspp_dropout=0.5,
        backbone_norm_layer=None,
        backbone_lr_mult=0.1,
        norm_layer=nn.BatchNorm2d,
        **kwargs
    ):
        super().__init__(norm_layer=norm_layer, **kwargs)

        self.feature_extractor = DeepLabV3Plus(
            backbone=backbone,
            ch=deeplab_ch,
            project_dropout=aspp_dropout,
            norm_layer=norm_layer,
            backbone_norm_layer=backbone_norm_layer,
        )
        self.feature_extractor.backbone.apply(LRMult(backbone_lr_mult))
        self.head = SepConvHead(
            1,
            in_channels=deeplab_ch,
            mid_channels=deeplab_ch // 2,
            num_layers=2,
            norm_layer=norm_layer,
        )

    def backbone_forward(self, image, coord_features=None):
        backbone_features = self.feature_extractor(image, coord_features)

        return {"instances": self.head(backbone_features[0])}