# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_model.ipynb.

# %% auto 0
__all__ = ['MODEL_DEFS', 'get_seg_model', 'get_preprocessiong_fn', 'get_model_name', 'calc_confusion_matrix',
           'display_confusion_matrix', 'save_model_checkpoint', 'create_checkpoint_data', 'restore_model']

# %% ../nbs/03_model.ipynb 3
import segmentation_models_pytorch as smp
import torch
from torchmetrics.classification import MulticlassConfusionMatrix
import matplotlib.pyplot as plt
import torch.nn.functional as F

# %% ../nbs/03_model.ipynb 5
def get_seg_model(model_def, classes, weights_fn=''):
    """
    Create a segmentation model, and load provided weighst
    """
    ENCODER = model_def['encoder']
    ENCODER_WEIGHTS = model_def['weights']
    ACTIVATION = None
    CLASSES = classes
    model_func = eval(model_def['arch'])
    model = model_func(
        encoder_name = ENCODER,
        encoder_weights=ENCODER_WEIGHTS,
        classes = len(CLASSES),
        activation = ACTIVATION
        )
    if weights_fn != '':
        checkpoint = torch.load(weights_fn)
        model.load_state_dict(checkpoint['model'])
    
    return model

# %% ../nbs/03_model.ipynb 7
def get_preprocessiong_fn(model_def):
    """
    Return the preprocessing function for a specific segmentation model
    """
    ENCODER = model_def['encoder']
    ENCODER_WEIGHTS = model_def['weights']
    return smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

def get_model_name(model_def, run_name=''):
    """
    Return the model name incl. a run name
    """
    arch_name = model_def['arch'].split('.')[-1]
    model_name = arch_name + '_' + model_def['encoder'] + run_name 
    return model_name

# %% ../nbs/03_model.ipynb 11
def calc_confusion_matrix(model, dl, num_classes):
    model.eval()
    metric = MulticlassConfusionMatrix(num_classes=num_classes,normalize='true')
    val = torch.zeros([num_classes,num_classes])
    size = 0
    for b in iter(dl):
        logits = model.predict(b[0])
        preds = torch.argmax(F.softmax(logits, dim=1), dim=1)
        val += metric(preds.to('cpu'), b[1].to('cpu')) * b[0].shape[0]
        size += b[0].shape[0]
    return val/size

def display_confusion_matrix(val, select_classes):
    fig, ax = plt.subplots(1,1, figsize=(7,7))
    MulticlassConfusionMatrix(num_classes=val.shape[0]).plot((val.round(decimals=2)*100.).trunc(),labels=select_classes, ax = ax)

# %% ../nbs/03_model.ipynb 13
def save_model_checkpoint(model, opt, path, **kwargs):
    d = {'model_state_dict' : model.state_dict(), 'optimizer_state_dict' : opt.state_dict()}
    d |= kwargs
    torch.save(d, path)

def create_checkpoint_data(model_def, select_classes, epoch, loss, metric):
    checkpoint_data = { 'arch' : model_def['arch'], 
                       'encoder' : model_def['encoder'],
                       'weights' : model_def['weights'],
                       'select_classes' : select_classes,
                       'epoch' : epoch,
                       'loss' : loss,
                       'metric' : metric
                        }
    return checkpoint_data

def restore_model(path):
    d = torch.load(path, map_location=torch.device('cpu') )
    model_def= {'arch': d['arch'], 'encoder' : d['encoder'], 'weights' : d['weights']}
    model = get_seg_model(model_def, d['select_classes'])
    model.load_state_dict(d['model_state_dict'])
    return model.eval(), d['select_classes'], model_def

# %% ../nbs/03_model.ipynb 15
MODEL_DEFS = []
MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'efficientnet-b5', 'weights': 'imagenet'}) #5 
MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'efficientnet-b7', 'weights': 'imagenet'}) #8
MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'efficientnet-b7', 'weights': 'imagenet'}) #9
MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'efficientnet-b5', 'weights': 'imagenet'}) #9
MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'efficientnet-b5', 'weights': 'imagenet'}) #10
MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'efficientnet-b7', 'weights': 'imagenet'}) #10
MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'xception', 'weights': 'imagenet'}) #11
#MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'resnext50_32x4d', 'weights': 'imagenet'}) #0
#MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'resnext101_32x8d', 'weights': 'imagenet'})#1
#MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'resnet101', 'weights': 'imagenet'})       #2
#MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'resnext50_32x4d', 'weights': 'ssl'}) #3
#MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'resnext101_32x8d', 'weights': 'imagenet'}) #4 
#MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'efficientnet-b6', 'weights': 'imagenet'}) #6 
#MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'resnext101_32x8d', 'weights': 'imagenet'}) #7 
#MODEL_DEFS.append({'arch': 'smp.Unet', 'encoder': 'inceptionresnetv2', 'weights': 'imagenet'}) #12
#MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'inceptionv4', 'weights': 'imagenet'}) #13
#MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'inceptionresnetv2', 'weights': 'imagenet'}) #14
#MODEL_DEFS.append({'arch': 'smp.UnetPlusPlus', 'encoder': 'xception', 'weights': 'imagenet'}) #15
#MODEL_DEFS.append({'arch': 'smp.DeepLabV3Plus', 'encoder': 'xception', 'weights': 'imagenet'}) #16
