Certainly! Below is a comprehensive Python script that integrates the DeepLabV3+ architecture with a MobileViT backbone. This script combines the definitions of both networks, modifies the MobileViT backbone to output the necessary feature maps for DeepLabV3+, and demonstrates how to instantiate and test the combined model.

Overview

1. MobileViT Backbone Modification:

   • The MobileViT class is adjusted to return two feature maps:
     • low_level_features: Extracted from an intermediate layer (after the second MV2Block).
     • x: The final deep features used by the ASPP module.

2. DeepLabV3+ Integration:

   • The DeepLab class is updated to utilize the modified MobileViT as its backbone.
   • The low_level_channels parameter is set appropriately based on the MobileViT architecture.

3. Complete Code Structure:

   • All necessary modules, helper functions, and classes are included in a single script.
   • A test case at the end verifies the integration by passing a sample input through the combined model.

Complete Code

python
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from einops import rearrange

# ---------------------------
# MobileViT Backbone
# ---------------------------

def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.SiLU()
    )


def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
    return nn.Sequential(
        nn.Conv2d(inp, oup, kernel_size, stride, kernel_size // 2, bias=False),
        nn.BatchNorm2d(oup),
        nn.SiLU()
    )


class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim=-1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        # x: (batch, patch, num_tokens, dim)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale  # (b, p, h, n, n)
        attn = self.attend(dots)  # (b, p, h, n, n)
        out = torch.matmul(attn, v)  # (b, p, h, n, d)
        out = rearrange(out, 'b p h n d -> b p n (h d)')
        return self.to_out(out)


class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
            ]))

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x


class MV2Block(nn.Module):
    def __init__(self, inp, oup, stride=1, expansion=4):
        super().__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(inp * expansion)
        self.use_res_connect = self.stride == 1 and inp == oup

        if expansion == 1:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileViTBlock(nn.Module):
    def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
        super().__init__()
        self.ph, self.pw = patch_size

        self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
        self.conv2 = conv_1x1_bn(channel, dim)

        self.transformer = Transformer(dim, depth, heads=4, dim_head=8, mlp_dim=mlp_dim, dropout=dropout)

        self.conv3 = conv_1x1_bn(dim, channel)
        self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)

    def forward(self, x):
        y = x.clone()

        # Local representations
        x = self.conv1(x)
        x = self.conv2(x)

        # Global representations
        _, _, h, w = x.shape
        x = rearrange(x, 'b c (h ph) (w pw) -> b (ph pw) (h w) c', ph=self.ph, pw=self.pw)
        x = self.transformer(x)
        x = rearrange(x, 'b (ph pw) (h w) c -> b c (h ph) (w pw)', h=h // self.ph, w=w // self.pw, ph=self.ph,
                      pw=self.pw)

        # Fusion
        x = self.conv3(x)
        x = torch.cat((x, y), 1)
        x = self.conv4(x)
        return x


class MobileViT(nn.Module):
    def __init__(self, image_size, dims, channels, num_classes, expansion=4, kernel_size=3, patch_size=(2, 2)):
        super().__init__()
        ih, iw = image_size
        ph, pw = patch_size
        assert ih % ph == 0 and iw % pw == 0

        L = [2, 4, 3]

        self.conv1 = conv_nxn_bn(3, channels[0], stride=2)

        self.mv2 = nn.ModuleList([])
        self.mv2.append(MV2Block(channels[0], channels[1], 1, expansion))
        self.mv2.append(MV2Block(channels[1], channels[2], 2, expansion))
        self.mv2.append(MV2Block(channels[2], channels[3], 1, expansion))
        self.mv2.append(MV2Block(channels[3], channels[4], 2, expansion))
        self.mv2.append(MV2Block(channels[4], channels[5], 1, expansion))
        self.mv2.append(MV2Block(channels[5], channels[6], 2, expansion))
        self.mv2.append(MV2Block(channels[6], channels[7], 1, expansion))
        self.mv2.append(MV2Block(channels[7], channels[8], 2, expansion))
        self.mv2.append(MV2Block(channels[8], channels[9], 1, expansion))
        self.mv2.append(MV2Block(channels[9], channels[10], 2, expansion))

        self.mvit = nn.ModuleList([])
        self.mvit.append(MobileViTBlock(dims[0], L[0], channels[5], kernel_size, patch_size, int(dims[0] * 2)))
        self.mvit.append(MobileViTBlock(dims[1], L[1], channels[7], kernel_size, patch_size, int(dims[1] * 4)))
        self.mvit.append(MobileViTBlock(dims[2], L[2], channels[9], kernel_size, patch_size, int(dims[2] * 4)))

        self.conv2 = conv_1x1_bn(channels[-2], channels[-1])
        self.conv3 = conv_1x1_bn(channels[-1], num_classes)

        self.upsample = nn.Upsample(size=image_size, mode='bilinear', align_corners=False)

        # Note: Global average pooling and classification layers are removed for backbone purposes

    def forward(self, x):
        x = self.conv1(x)          # Output: channels[0], stride=2
        x = self.mv2[0](x)         # Output: channels[1]
        x = self.mv2[1](x)         # Output: channels[2], stride=2

        # Extract low_level_features here (after second MV2Block)
        low_level_features = x.clone()

        x = self.mv2[2](x)         # Output: channels[3]
        x = self.mv2[3](x)         # Output: channels[4], stride=2
        x = self.mv2[4](x)         # Output: channels[5]
        x = self.mvit[0](x)        # MobileViTBlock 1
        x = self.mv2[5](x)         # Output: channels[6], stride=2
        x = self.mvit[1](x)        # MobileViTBlock 2
        x = self.mv2[6](x)         # Output: channels[7]
        x = self.mvit[2](x)        # MobileViTBlock 3
        x = self.mv2[7](x)         # Output: channels[8], stride=2
        x = self.mv2[8](x)         # Output: channels[9]
        x = self.mv2[9](x)         # Output: channels[10], stride=2
        x = self.conv2(x)          # Output: channels[10]

        # For backbone purposes, we return low_level_features and x
        return low_level_features, x


# ---------------------------
# ASPP Module
# ---------------------------

class ASPP(nn.Module):
    def __init__(self, dim_in, dim_out, rate=1, bn_mom=0.1):
        super(ASPP, self).__init__()
        self.branch1 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 1, 1, padding=0, dilation=rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch2 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=6 * rate, dilation=6 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch3 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=12 * rate, dilation=12 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch4 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=18 * rate, dilation=18 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch5_conv = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=True)
        self.branch5_bn = nn.BatchNorm2d(dim_out, momentum=bn_mom)
        self.branch5_relu = nn.ReLU(inplace=True)

        self.conv_cat = nn.Sequential(
            nn.Conv2d(dim_out * 5, dim_out, 1, 1, padding=0, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        [b, c, row, col] = x.size()
        # Five branches
        conv1x1 = self.branch1(x)
        conv3x3_1 = self.branch2(x)
        conv3x3_2 = self.branch3(x)
        conv3x3_3 = self.branch4(x)
        # Fifth branch: global average pooling + conv
        global_feature = torch.mean(x, 2, True)
        global_feature = torch.mean(global_feature, 3, True)
        global_feature = self.branch5_conv(global_feature)
        global_feature = self.branch5_bn(global_feature)
        global_feature = self.branch5_relu(global_feature)
        global_feature = F.interpolate(global_feature, (row, col), None, 'bilinear', True)

        # Concatenate and project
        feature_cat = torch.cat([conv1x1, conv3x3_1, conv3x3_2, conv3x3_3, global_feature], dim=1)
        result = self.conv_cat(feature_cat)
        return result


# ---------------------------
# DeepLabV3+ with MobileViT Backbone
# ---------------------------

class DeepLab(nn.Module):
    def __init__(self, num_classes, backbone="mobilivit", pretrained_backbone=True, downsample_factor=16):
        super(DeepLab, self).__init__()
        if backbone == "xception":
            raise NotImplementedError("Xception backbone is not implemented in this script.")
        elif backbone == "mobilivit":
            # Initialize MobileViT backbone
            # Define MobileViT parameters as per your requirements
            dims = [96, 120, 144]
            channels = [16, 32, 48, 48, 64, 64, 80, 80, 96, 96, 384]
            self.backbone = MobileViT(
                image_size=(256, 256),
                dims=dims,
                channels=channels,
                num_classes=1000,  # Not used in backbone
                expansion=2,
                kernel_size=3,
                patch_size=(2, 2)
            )
            in_channels = channels[-1]  # 384
            low_level_channels = channels[2]  # 48
        else:
            raise ValueError('Unsupported backbone - `{}`, Use mobilivit.'.format(backbone))

        # ASPP Module
        self.aspp = ASPP(dim_in=in_channels, dim_out=256, rate=16 // downsample_factor)

        # Shortcut convolution for low-level features
        self.shortcut_conv = nn.Sequential(
            nn.Conv2d(low_level_channels, 48, 1),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )

        # Convolution layers after concatenation
        self.cat_conv = nn.Sequential(
            nn.Conv2d(48 + 256, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),

            nn.Conv2d(256, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.Dropout(0.1),
        )
        self.cls_conv = nn.Conv2d(256, num_classes, 1, stride=1)

    def forward(self, x):
        H, W = x.size(2), x.size(3)
        # Backbone forward
        low_level_features, x = self.backbone(x)
        # ASPP
        x = self.aspp(x)
        # Process low-level features
        low_level_features = self.shortcut_conv(low_level_features)

        # Upsample ASPP output to match low-level features
        x = F.interpolate(x, size=low_level_features.size()[2:], mode='bilinear', align_corners=True)
        # Concatenate
        x = self.cat_conv(torch.cat((x, low_level_features), dim=1))
        # Classification layer
        x = self.cls_conv(x)
        # Final upsampling to input size
        x = F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True)
        return x


# ---------------------------
# Parameter Counting Utility
# ---------------------------

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# ---------------------------
# Testing the Combined Model
# ---------------------------

if __name__ == '__main__':
    # Example input
    img = torch.randn(1, 3, 256, 256)

    # Initialize DeepLab with MobileViT backbone
    model = DeepLab(
        num_classes=3,               # Example for 3 classes
        backbone="mobilivit",        # Specify the backbone
        pretrained_backbone=True,    # Whether to use pretrained weights (if available)
        downsample_factor=16         # Downsample factor
    )

    # Forward pass
    out = model(img)
    print("Output shape:", out.shape)  # Expected: (1, 3, 256, 256)

    # Optional: Print number of parameters
    print("Total trainable parameters:", count_parameters(model))

Explanation of the Integration

1. MobileViT Backbone Modification:

   • Feature Extraction: The MobileViT class is designed to extract features at different stages. For DeepLabV3+, we need two outputs:
     • low_level_features: Captured after the second MV2Block (channels[2] = 48).
     • x: The final deep features after all MV2Block and MobileViTBlock layers (channels[-1] = 384).
   • Forward Method: The forward method of MobileViT is modified to return both low_level_features and x.

2. DeepLabV3+ Configuration:

   • Backbone Selection: The DeepLab class is configured to use the MobileViT backbone by specifying backbone="mobilivit".
   • ASPP Module: The ASPP module is initialized with dim_in set to the number of channels from the backbone's output (384 in this case) and dim_out=256.
   • Shortcut Convolution: Processes the low_level_features to reduce channel dimensions to 48.
   • Concatenation and Classification: The ASPP output and processed low_level_features are concatenated and passed through additional convolutional layers to produce the final segmentation map.

3. Testing:

   • An example input tensor of shape (1, 3, 256, 256) is passed through the combined model.
   • The output shape is printed to verify it matches the input spatial dimensions with the desired number of classes (3 in this example).
   • The total number of trainable parameters is also printed for reference.

Notes

• Pretrained Weights: The pretrained_backbone=True parameter is included for scenarios where pretrained weights for MobileViT are available. Ensure that you load appropriate pretrained weights if necessary.

• Backbone Flexibility: While this script focuses on integrating MobileViT as the backbone, the structure allows for easy extension to other backbones (e.g., Xception) by following a similar pattern.

• Adjustable Parameters: You can modify parameters such as num_classes, image_size, and downsample_factor based on your specific requirements.

This integrated model leverages the strengths of MobileViT for efficient feature extraction and the powerful segmentation capabilities of DeepLabV3+. It is well-suited for tasks requiring accurate and efficient semantic segmentation.

The error you're encountering stems from a mismatch between the expected number of input channels for a convolutional layer and the actual number of channels being passed during the forward pass of your MobileViT backbone within the DeepLabV3+ architecture.

Error Breakdown

Error Message:

RuntimeError: Given groups=1, weight of size [96, 96, 3, 3], expected input[1, 80, 16, 16] to have 96 channels, but got 80 channels instead

Cause:

• The MobileViTBlock is expecting an input with 96 channels, but it is receiving an input with 80 channels.
• This discrepancy occurs because the MobileViT backbone's forward pass is incorrectly ordering the application of MobileViTBlock and MV2Block layers, leading to a mismatch in channel dimensions.

Detailed Analysis

Let's delve into the MobileViT class to understand the flow of channels:

1. Channel Configuration:

   python
   channels = [16, 32, 48, 48, 64, 64, 80, 80, 96, 96, 384]

2. Layer Initialization:

   • MV2Blocks:
     • self.mv2 contains 10 MV2Block instances, transitioning through the defined channels.
   • MobileViTBlocks:
     • self.mvit contains 3 MobileViTBlock instances:
       • self.mvit[0] expects channels[5] = 64 channels.
       • self.mvit[1] expects channels[7] = 80 channels.
       • self.mvit[2] expects channels[9] = 96 channels.

3. Forward Pass Issue:

   • In the original forward pass:

     python
     x = self.mv2[6](x)         # Output channels[7] = 80
     x = self.mvit[2](x)        # Expects 96 channels

   • Here, self.mvit[2] is being applied to an input with 80 channels instead of the expected 96 channels, causing the runtime error.

Solution

Rearrange the Forward Pass Order:

To ensure that self.mvit[2] receives the correct number of input channels (96 channels), you need to adjust the order in which layers are applied in the forward method of the MobileViT class.

Correct Order:

1. Apply self.mv2[6] to transition from 80 to 80 channels.
2. Apply self.mv2[7] to transition from 80 to 96 channels.
3. Now, self.mvit[2] can safely process the input with 96 channels.

Updated Forward Method:

python
def forward(self, x):
    x = self.conv1(x)          # Output: channels[0]=16
    x = self.mv2[0](x)         # Output: channels[1]=32
    x = self.mv2[1](x)         # Output: channels[2]=48, stride=2

    low_level_features = x.clone() # channels[2]=48

    x = self.mv2[2](x)         # Output: channels[3]=48
    x = self.mv2[3](x)         # Output: channels[4]=64, stride=2
    x = self.mv2[4](x)         # Output: channels[5]=64
    x = self.mvit[0](x)        # MobileViTBlock 1: Input channels=64, Output channels=64

    x = self.mv2[5](x)         # Output: channels[6]=80, stride=2
    x = self.mvit[1](x)        # MobileViTBlock 2: Input channels=80, Output channels=80

    x = self.mv2[6](x)         # Output: channels[7]=80
    x = self.mv2[7](x)         # Output: channels[8]=96
    x = self.mvit[2](x)        # MobileViTBlock 3: Input channels=96, Output channels=96

    x = self.mv2[8](x)         # Output: channels[9]=96
    x = self.mv2[9](x)         # Output: channels[10]=384
    x = self.conv2(x)          # Output: channels[10]=384

    return low_level_features, x

Complete Corrected Code

Below is the corrected and comprehensive Python script that integrates the DeepLabV3+ architecture with a MobileViT backbone. The primary change is the rearrangement of the forward pass in the MobileViT class to ensure channel compatibility.

python
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from einops import rearrange

# ---------------------------
# MobileViT Backbone
# ---------------------------

def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.SiLU()
    )


def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
    return nn.Sequential(
        nn.Conv2d(inp, oup, kernel_size, stride, kernel_size // 2, bias=False),
        nn.BatchNorm2d(oup),
        nn.SiLU()
    )


class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim=-1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        # x: (batch, patch, num_tokens, dim)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale  # (b, p, h, n, n)
        attn = self.attend(dots)  # (b, p, h, n, n)
        out = torch.matmul(attn, v)  # (b, p, h, n, d)
        out = rearrange(out, 'b p h n d -> b p n (h d)')
        return self.to_out(out)


class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
            ]))

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x


class MV2Block(nn.Module):
    def __init__(self, inp, oup, stride=1, expansion=4):
        super().__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(inp * expansion)
        self.use_res_connect = self.stride == 1 and inp == oup

        if expansion == 1:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.SiLU(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileViTBlock(nn.Module):
    def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
        super().__init__()
        self.ph, self.pw = patch_size

        self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
        self.conv2 = conv_1x1_bn(channel, dim)

        self.transformer = Transformer(dim, depth, heads=4, dim_head=8, mlp_dim=mlp_dim, dropout=dropout)

        self.conv3 = conv_1x1_bn(dim, channel)
        self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)

    def forward(self, x):
        y = x.clone()

        # Local representations
        x = self.conv1(x)
        x = self.conv2(x)

        # Global representations
        _, _, h, w = x.shape
        x = rearrange(x, 'b c (h ph) (w pw) -> b (ph pw) (h w) c', ph=self.ph, pw=self.pw)
        x = self.transformer(x)
        x = rearrange(x, 'b (ph pw) (h w) c -> b c (h ph) (w pw)', h=h // self.ph, w=w // self.pw, ph=self.ph,
                      pw=self.pw)

        # Fusion
        x = self.conv3(x)
        x = torch.cat((x, y), 1)
        x = self.conv4(x)
        return x


class MobileViT(nn.Module):
    def __init__(self, image_size, dims, channels, num_classes, expansion=4, kernel_size=3, patch_size=(2, 2)):
        super().__init__()
        ih, iw = image_size
        ph, pw = patch_size
        assert ih % ph == 0 and iw % pw == 0

        L = [2, 4, 3]

        self.conv1 = conv_nxn_bn(3, channels[0], stride=2)

        self.mv2 = nn.ModuleList([])
        self.mv2.append(MV2Block(channels[0], channels[1], 1, expansion))
        self.mv2.append(MV2Block(channels[1], channels[2], 2, expansion))
        self.mv2.append(MV2Block(channels[2], channels[3], 1, expansion))
        self.mv2.append(MV2Block(channels[3], channels[4], 2, expansion))
        self.mv2.append(MV2Block(channels[4], channels[5], 1, expansion))
        self.mv2.append(MV2Block(channels[5], channels[6], 2, expansion))
        self.mv2.append(MV2Block(channels[6], channels[7], 1, expansion))
        self.mv2.append(MV2Block(channels[7], channels[8], 2, expansion))
        self.mv2.append(MV2Block(channels[8], channels[9], 1, expansion))
        self.mv2.append(MV2Block(channels[9], channels[10], 2, expansion))

        self.mvit = nn.ModuleList([])
        self.mvit.append(MobileViTBlock(dims[0], L[0], channels[5], kernel_size, patch_size, int(dims[0] * 2)))
        self.mvit.append(MobileViTBlock(dims[1], L[1], channels[7], kernel_size, patch_size, int(dims[1] * 4)))
        self.mvit.append(MobileViTBlock(dims[2], L[2], channels[9], kernel_size, patch_size, int(dims[2] * 4)))

        self.conv2 = conv_1x1_bn(channels[-2], channels[-1])
        self.conv3 = conv_1x1_bn(channels[-1], num_classes)

        self.upsample = nn.Upsample(size=image_size, mode='bilinear', align_corners=False)

        # Note: Global average pooling and classification layers are removed for backbone purposes

    def forward(self, x):
        x = self.conv1(x)          # Output: channels[0]=16
        x = self.mv2[0](x)         # Output: channels[1]=32
        x = self.mv2[1](x)         # Output: channels[2]=48, stride=2

        low_level_features = x.clone() # channels[2]=48

        x = self.mv2[2](x)         # Output: channels[3]=48
        x = self.mv2[3](x)         # Output: channels[4]=64, stride=2
        x = self.mv2[4](x)         # Output: channels[5]=64
        x = self.mvit[0](x)        # MobileViTBlock 1: Input channels=64, Output channels=64

        x = self.mv2[5](x)         # Output: channels[6]=80, stride=2
        x = self.mvit[1](x)        # MobileViTBlock 2: Input channels=80, Output channels=80

        x = self.mv2[6](x)         # Output: channels[7]=80
        x = self.mv2[7](x)         # Output: channels[8]=96
        x = self.mvit[2](x)        # MobileViTBlock 3: Input channels=96, Output channels=96

        x = self.mv2[8](x)         # Output: channels[9]=96
        x = self.mv2[9](x)         # Output: channels[10]=384
        x = self.conv2(x)          # Output: channels[10]=384

        return low_level_features, x

# ---------------------------
# ASPP Module
# ---------------------------

class ASPP(nn.Module):
    def __init__(self, dim_in, dim_out, rate=1, bn_mom=0.1):
        super(ASPP, self).__init__()
        self.branch1 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 1, 1, padding=0, dilation=rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch2 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=6 * rate, dilation=6 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch3 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=12 * rate, dilation=12 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch4 = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, 3, 1, padding=18 * rate, dilation=18 * rate, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )
        self.branch5_conv = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=True)
        self.branch5_bn = nn.BatchNorm2d(dim_out, momentum=bn_mom)
        self.branch5_relu = nn.ReLU(inplace=True)

        self.conv_cat = nn.Sequential(
            nn.Conv2d(dim_out * 5, dim_out, 1, 1, padding=0, bias=True),
            nn.BatchNorm2d(dim_out, momentum=bn_mom),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        [b, c, row, col] = x.size()
        # Five branches
        conv1x1 = self.branch1(x)
        conv3x3_1 = self.branch2(x)
        conv3x3_2 = self.branch3(x)
        conv3x3_3 = self.branch4(x)
        # Fifth branch: global average pooling + conv
        global_feature = torch.mean(x, 2, True)
        global_feature = torch.mean(global_feature, 3, True)
        global_feature = self.branch5_conv(global_feature)
        global_feature = self.branch5_bn(global_feature)
        global_feature = self.branch5_relu(global_feature)
        global_feature = F.interpolate(global_feature, (row, col), None, 'bilinear', True)

        # Concatenate and project
        feature_cat = torch.cat([conv1x1, conv3x3_1, conv3x3_2, conv3x3_3, global_feature], dim=1)
        result = self.conv_cat(feature_cat)
        return result

# ---------------------------
# DeepLabV3+ with MobileViT Backbone
# ---------------------------

class DeepLab(nn.Module):
    def __init__(self, num_classes, backbone="mobilivit", pretrained_backbone=True, downsample_factor=16):
        super(DeepLab, self).__init__()
        if backbone == "xception":
            raise NotImplementedError("Xception backbone is not implemented in this script.")
        elif backbone == "mobilivit":
            # Initialize MobileViT backbone
            # Define MobileViT parameters as per your requirements
            dims = [96, 120, 144]
            channels = [16, 32, 48, 48, 64, 64, 80, 80, 96, 96, 384]
            self.backbone = MobileViT(
                image_size=(256, 256),
                dims=dims,
                channels=channels,
                num_classes=1000,  # Not used in backbone
                expansion=2,
                kernel_size=3,
                patch_size=(2, 2)
            )
            in_channels = channels[-1]  # 384
            low_level_channels = channels[2]  # 48
        else:
            raise ValueError('Unsupported backbone - `{}`, Use mobilivit.'.format(backbone))

        # ASPP Module
        self.aspp = ASPP(dim_in=in_channels, dim_out=256, rate=16 // downsample_factor)

        # Shortcut convolution for low-level features
        self.shortcut_conv = nn.Sequential(
            nn.Conv2d(low_level_channels, 48, 1),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )

        # Convolution layers after concatenation
        self.cat_conv = nn.Sequential(
            nn.Conv2d(48 + 256, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),

            nn.Conv2d(256, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.Dropout(0.1),
        )
        self.cls_conv = nn.Conv2d(256, num_classes, 1, stride=1)

    def forward(self, x):
        H, W = x.size(2), x.size(3)
        # Backbone forward
        low_level_features, x = self.backbone(x)
        # ASPP
        x = self.aspp(x)
        # Process low-level features
        low_level_features = self.shortcut_conv(low_level_features)

        # Upsample ASPP output to match low-level features
        x = F.interpolate(x, size=low_level_features.size()[2:], mode='bilinear', align_corners=True)
        # Concatenate
        x = self.cat_conv(torch.cat((x, low_level_features), dim=1))
        # Classification layer
        x = self.cls_conv(x)
        # Final upsampling to input size
        x = F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True)
        return x

# ---------------------------
# Parameter Counting Utility
# ---------------------------

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# ---------------------------
# Testing the Combined Model
# ---------------------------

if __name__ == '__main__':
    # Example input
    img = torch.randn(1, 3, 256, 256)

    # Initialize DeepLab with MobileViT backbone
    model = DeepLab(
        num_classes=3,               # Example for 3 classes
        backbone="mobilivit",        # Specify the backbone
        pretrained_backbone=True,    # Whether to use pretrained weights (if available)
        downsample_factor=16         # Downsample factor
    )

    # Forward pass
    out = model(img)
    print("Output shape:", out.shape)  # Expected: (1, 3, 256, 256)

    # Optional: Print number of parameters
    print("Total trainable parameters:", count_parameters(model))