import numpy as np
import megengine
import megengine.module as M
import megengine.functional as F
from edit.models.common import ShuffleV2Block, CoordAtt
import math
from . import default_init_weights


class MobileNeXt(M.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        """
            默认使用coordinate attention在第一个dwise之后
            https://github.com/Andrew-Qibin/CoordAttention/blob/main/coordatt.py
        """
        super(MobileNeXt, self).__init__()
        self.dconv1 = M.ConvRelu2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=(kernel_size//2), groups=in_channels)
        self.CA = CoordAtt(inp = out_channels, oup=out_channels)
        self.conv1 = 
M.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, padding=0)