# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00vision_models.ipynb (unless otherwise specified).

__all__ = ['gem', 'GeM', 'AdaptiveConcatPool2d_GeM', 'MobileNetV2', 'ResNet_', 'ResNet18_swsl', 'ResNet50_swsl',
           'ResNet50_32x4d_swsl', 'xResNet_', 'xResNet50_ssa', 'ResNeSt_', 'ResNeSt50', 'ResNeSt101', 'ResNeSt200',
           'ResNeSt269', 'ResNeSt50_fast_1s1x64d', 'ResNeSt50_fast_1s2x40d', 'ResNeSt50_fast_1s4x24d',
           'ResNeSt50_fast_2s1x64d', 'ResNeSt50_fast_2s2x40d', 'ResNeSt50_fast_4s1x64d', 'ResNeSt50_fast_4s2x40d',
           'DenseNet_', 'DenseNet121', 'DenseNet169', 'DenseNet201', 'DenseNet161', 'DenseNetBlur121', 'EfficientNet_',
           'EfficientNetB0', 'EfficientNetB1', 'EfficientNetB2', 'EfficientNetB3', 'EfficientNetB4', 'EfficientNetB5',
           'EfficientNetB6', 'EfficientNetB7', 'Head', 'EmbResNeSt_', 'EmbResNeSt50', 'get_model']

# Cell
from fastai.vision.all import *
from efficientnet_pytorch.utils import Conv2dStaticSamePadding
from efficientnet_pytorch import EfficientNet
import timm
from ..external import xresnet_ssa

# Cell
def gem(x, p, eps=1e-6):
    t = x.type()
    x = x.float()
    x = F.avg_pool2d(x.clamp(min=eps).pow(p.view(1,-1,1,1)),
                     (x.size(-2), x.size(-1))).pow(1.0/p.view(1,-1,1,1))
    x = x.type(t)
    return x

class GeM(nn.Module):
    def __init__(self, nc=1, p=3, eps=1e-6):
        super(GeM,self).__init__()
        self.p = nn.Parameter(torch.ones(nc)*p)
        self.eps = eps

    def forward(self, x):
        return gem(x, p=self.p, eps=self.eps)

    def __repr__(self):
        return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'\
        .format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')'

#export
class AdaptiveConcatPool2d_GeM(Module):
    "Layer that concats `AdaptiveAvgPool2d` and `AdaptiveMaxPool2d`"
    def __init__(self, size=None):
        self.size = size or 1
        self.ap = nn.AdaptiveAvgPool2d(self.size)
        self.mp = nn.AdaptiveMaxPool2d(self.size)
        self.gem = GeM()
    def forward(self, x): return torch.cat([self.mp(x), self.ap(x), self.gem(x)], 1)

# Cell
class MobileNetV2(Module):
    def __init__(self, num_classes, head_ps=0.5, pretrained=True, in_channels=3, **kwargs):
        body = torch.hub.load('pytorch/vision:v0.6.0', 'mobilenet_v2', pretrained=pretrained)
        body = nn.Sequential(*children_and_parameters(body)[:-1])
        num_features = 1280
        model = nn.Sequential(body, create_head(num_features*2, num_classes, ps=head_ps))
        if in_channels == 1:
            out_ch =model[0][0][0][0].out_channels
            k_sz = model[0][0][0][0].kernel_size
            stride = model[0][0][0][0].stride
            pad = model[0][0][0][0].padding
            w = model[0][0][0][0].weight.data[:,0].unsqueeze(1)
            model[0][0][0][0] = nn.Conv2d(1, out_ch, kernel_size=k_sz, stride=stride, padding=pad, bias=False)
            model[0][0][0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

# Cell
class ResNet_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.5, pretrained=True, in_channels=3, **kwargs):
        body = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models', 'resnet18_swsl')
        body = nn.Sequential(*children_and_parameters(body)[:-2])
        num_features = body[-1][-1].bn2.num_features
        model = nn.Sequential(body, create_head(num_features*2, num_classes, ps=head_ps))
        if in_channels == 1:
            out_ch = model[0][0].out_channels
            k_sz = model[0][0].kernel_size
            stride = model[0][0].stride
            pad = model[0][0].padding
            w = model[0][0].weight.data[:,0].unsqueeze(1)
            model[0][0] = nn.Conv2d(1, out_ch, kernel_size=k_sz, stride=stride, padding=pad, bias=False)
            model[0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class ResNet18_swsl(ResNet_):
    arch='resnet18_swsl'

class ResNet50_swsl(ResNet_):
    arch='resnet50_swsl'

class ResNet50_32x4d_swsl(ResNet_):
    arch='resnext50_32x4d_swsl'

# Cell
class xResNet_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.8, pretrained=True, in_channels=3, **kwargs):
        model = eval(f'xresnet_ssa.{self.arch}')(pretrained=pretrained)
        model = nn.Sequential(*[*children_and_parameters(model)[:-2],
                *children_and_parameters(create_head(2048*2, num_classes, ps=head_ps))])
        if in_channels == 1:
            w = model[0][0].weight.data[:,0].unsqueeze(1)
            model[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2),
                                    padding=(1, 1), bias=False)
            model[0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class xResNet50_ssa(xResNet_):
    arch='xresnet50'

# Cell
class ResNeSt_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.8, pretrained=True, in_channels=3, **kwargs):
        try:
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        except:
            torch.hub.list('zhanghang1989/ResNeSt', force_reload=True)
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        model = nn.Sequential(*[*children_and_parameters(model)[:-2],
                *children_and_parameters(create_head(2048*2, num_classes, ps=head_ps))])
        if in_channels == 1:
            w = model[0][0].weight.data[:,0].unsqueeze(1)
            model[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2),
                                    padding=(1, 1), bias=False)
            model[0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class ResNeSt50(ResNeSt_):
    arch='resnest50'

class ResNeSt101(ResNeSt_):
    arch='resnest101'

class ResNeSt200(ResNeSt_):
    arch='resnest200'

class ResNeSt269(ResNeSt_):
    arch='resnest269'

class ResNeSt50_fast_1s1x64d(ResNeSt_):
    arch='resnest50_fast_1s1x64d'

class ResNeSt50_fast_1s2x40d(ResNeSt_):
    arch='resnest50_fast_1s2x40d'

class ResNeSt50_fast_1s4x24d(ResNeSt_):
    arch='resnest50_fast_1s4x24d'

class ResNeSt50_fast_2s1x64d(ResNeSt_):
    arch='resnest50_fast_2s1x64d'

class ResNeSt50_fast_2s2x40d(ResNeSt_):
    arch='resnest50_fast_2s2x40d'

class ResNeSt50_fast_4s1x64d(ResNeSt_):
    arch='resnest50_fast_4s1x64d'

class ResNeSt50_fast_4s2x40d(ResNeSt_):
    arch='resnest50_fast_4s2x40d'

# Cell
class DenseNet_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.5, pretrained=True, in_channels=3, **kwargs):
        body = eval(f'models.{self.arch}')(pretrained=pretrained).features
        num_features = children_and_parameters(body)[-1].num_features
        model = nn.Sequential(body, create_head(num_features*2, num_classes, ps=head_ps))
        if in_channels == 1:
            w = model[0].conv0.weight.data[:,0].unsqueeze(1)
            model[0].conv0 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2,2), padding=(3, 3), bias=False)
            model[0].conv0.weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class DenseNet121(DenseNet_):
    arch='densenet121'

class DenseNet169(DenseNet_):
    arch='densenet169'

class DenseNet201(DenseNet_):
    arch='densenet201'

class DenseNet161(DenseNet_):
    arch='densenet161'

# Cell
class DenseNetBlur121(Module):
    arch='densenetblur121d'
    def __init__(self, num_classes, head_ps=0.5, pretrained=True, in_channels=3, **kwargs):
        body =  timm.create_model(self.arch,num_classes=num_classes,
                                  pretrained=pretrained).features
        num_features = body[-1].num_features
        model = nn.Sequential(body, create_head(num_features*2, num_classes, ps=head_ps))
        if in_channels == 1:
            w = model[0].conv0.weight.data[:,0].unsqueeze(1)
            model[0].conv0 = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
            model[0].conv0.weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

# Cell
class EfficientNet_(Module):
    arch=None
    def __init__(self, num_classes, head_ps=0.5, pretrained=True, in_channels=3, **kwargs):
        self.effnet = EfficientNet.from_pretrained(self.arch)
        if in_channels == 1:
            w = self.effnet._conv_stem.weight.data[:,0].unsqueeze(1)
            out_channels = self.effnet._conv_stem.out_channels
            self.effnet._conv_stem = Conv2dStaticSamePadding(1, out_channels, kernel_size=(3,3), stride=(2,2),
                                                             bias=False, image_size=(128,128))
            self.effnet._conv_stem.weight.data = w
        num_features = self.effnet._bn1.num_features
        self.head = create_head(num_features*2, num_classes, ps=head_ps)
    def forward(self, x):
        x = self.effnet.extract_features(x)
        return self.head(x)

class EfficientNetB0(EfficientNet_):
    arch = 'efficientnet-b0'

class EfficientNetB1(EfficientNet_):
    arch = 'efficientnet-b1'

class EfficientNetB2(EfficientNet_):
    arch = 'efficientnet-b2'

class EfficientNetB3(EfficientNet_):
    arch = 'efficientnet-b3'

class EfficientNetB4(EfficientNet_):
    arch = 'efficientnet-b4'

class EfficientNetB5(EfficientNet_):
    arch = 'efficientnet-b5'

class EfficientNetB6(EfficientNet_):
    arch = 'efficientnet-b6'

class EfficientNetB7(EfficientNet_):
    arch = 'efficientnet-b7'

# Cell
class Head(nn.Module):
    def __init__(self, in_channels, emb_sz=128):
        super().__init__()
        self.flat = nn.Sequential(
            AdaptiveConcatPool2d(1))
        self.flatten = Flatten()
        self.bn0 = nn.BatchNorm1d(in_channels)
        self.lin0 = nn.Linear(in_channels, in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.bn1 = nn.BatchNorm1d(in_channels)
        self.lin1 = nn.Linear(in_channels, emb_sz)

    def forward(self, x):
        x = self.relu(self.lin0(self.bn0(x)))
        return self.lin1(self.bn1(x))

class EmbResNeSt_(Module):
    arch=None
    def __init__(self, emb_sz=128, head_ps=0.5, pretrained=True, **kwargs):
        try:
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        except:
            torch.hub.list('zhanghang1989/ResNeSt', force_reload=True)
            model = torch.hub.load('zhanghang1989/ResNeSt', self.arch, pretrained=pretrained)
        model = nn.Sequential(*[*children_and_parameters(model)[:-2],
                *children_and_parameters(Head(2048*2, emb_sz))])
        w = model[0][0].weight.data[:,0].unsqueeze(1)
        model[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2),
                                padding=(1, 1), bias=False)
        model[0][0].weight.data = w
        self.layers = model
    def forward(self, x):
        return self.layers(x)

class EmbResNeSt50(EmbResNeSt_):
    arch='resnest50'

# Cell
def get_model(arch, num_classes, head_ps=0.5, in_channels=3, pretrained=True,
              device=torch.device("cuda:0"), **kwargs):
    if arch == 'resnest50':
        return ResNeSt50(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                         in_channels=in_channels, **kwargs).to(device)
    if arch == 'resnest101':
        return ResNeSt101(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                           in_channels=in_channels, **kwargs).to(device)
    if arch == 'densenet121':
        return DenseNet121(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'densenet169':
        return DenseNet169(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'densenet201':
        return DenseNet201(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'densenet161':
        return DenseNet161(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'densenetblur121d':
        return DenseNetBlur121(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb0':
        return EfficientNetB0(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb1':
        return EfficientNetB1(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb2':
        return EfficientNetB2(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb3':
        return EfficientNetB3(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb4':
        return EfficientNetB4(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb5':
        return EfficientNetB5(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb6':
        return EfficientNetB6(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'efficientnetb7':
        return EfficientNetB7(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'resnet50_32x4d_swsl':
        return ResNet50_32x4d_swsl(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                               in_channels=in_channels, **kwargs).to(device)
    if arch == 'mobilenetv2':
        return MobileNetV2(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    if arch == 'embresnest50':
        return EmbResNeSt50(num_classes=num_classes, head_ps=head_ps, pretrained=pretrained,
                            in_channels=in_channels, **kwargs).to(device)
    raise Exception(f'{arch} not defined.')