import os
if not os.path.exists("/home/featurize/data/anno.zip"):
    !featurize dataset download 757181ec-0841-4468-b701-54ea53ad2d3c
if not os.path.exists("/home/featurize/data/forgery_round1_train_20220217.zip"):
    !featurize dataset download e4dc9a7b-ded4-453d-b887-1f3fb04a5f94

from IPython.display import clear_output
%pip install minetorch seaborn albumentations segmentation_models_pytorch
clear_output()

import os
import cv2
import torch
import random
import minetorch

import numpy as np
import pandas as pd
import seaborn as sns
import albumentations as albu
import matplotlib.pyplot as plt
import segmentation_models_pytorch as smp

from albumentations.pytorch import ToTensorV2
from albumentations import Normalize, Compose
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

CONFIG = {
    'DIR': 'Experiments',
    'img_size': 224,
    'model_name': 'timm-efficientnet-b0',
    'fold': 0,
    'lr': 1e-4,
    'batch_size': 128,
    'epoch': 20
}

class AliDataset(Dataset):

    def __init__(self, df, transforms, train=False):
        self.df = df
        self.transforms = transforms
        self.train = train
    
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        fn = row.image
        image = cv2.imread(os.path.join(row['image_path'], str(fn) + '.jpg'))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        masks = cv2.imread(os.path.join(row['mask_path'], str(fn) + '.png'), cv2.IMREAD_GRAYSCALE)/255

        augmented = self.transforms(image=image, mask=masks)
        img, mask = augmented['image'], augmented['mask']

        return img.float(), mask.float()

    def __len__(self):
        return len(self.df)

def create_transformation(phase):
    
    transforms = albu.Compose([
        albu.RandomBrightnessContrast(p=0.5),
        albu.OneOf([
            albu.Rotate(limit=[90,90], p=0.5),
            albu.Rotate(limit=[270,270], p=0.5),
        ], p=0.5),
        albu.RandomResizedCrop(CONFIG['img_size'], CONFIG['img_size'], scale=(0.8, 1.0), ratio=(0.75, 1.3333333333333333)),
        albu.Resize(CONFIG['img_size'], CONFIG['img_size'], p=1),
        ToTensorV2()
    ])
    
    if phase == 'val':
        transforms = albu.Compose([
            albu.Resize(CONFIG['img_size'], CONFIG['img_size'], p=1),
            ToTensorV2()
        ])
    
    return transforms

from minetorch.plugin import Plugin

def dice_coeff(pred, target):
    smooth = 1e-12
    num = pred.size(0)
    m1 = pred.view(num, -1).float()  # Flatten
    m2 = target.view(num, -1).float()  # Flatten
    intersection = (m1 * m2).sum().float()
    return (2. * intersection + smooth) / (m1.sum() + m2.sum() + smooth)

def iou_pytorch(outputs: torch.Tensor, labels: torch.Tensor):
    SMOOTH = 1e-12
    outputs = outputs.squeeze(1)  # BATCH x 1 x H x W => BATCH x H x W
    intersection = (outputs * labels).float().sum((1, 2))  
    union = (outputs + labels).float().sum((1, 2))
    iou = (intersection + SMOOTH) / (union - intersection + SMOOTH)
    return iou

class MultiClassesSegmentationMetricWithLogic(Plugin):
    """MultiClassesClassificationMetric
    This can be used directly if your loss function is torch.nn.CrossEntropy
    """

    def __init__(self,
                 miou=True,
                 mf1=True,
                 sheet_key_prefix=''):

        super().__init__(sheet_key_prefix)
        self.miou = miou
        self.mf1 = mf1


    def before_init(self):
        self.create_sheet_column('miou', 'miou')

    def before_epoch_start(self, epoch):
        self.iou = 0
        self.f1 = 0
        self.len = 0

    def after_val_iteration_ended(self, predicts, data, **ignore):

        predicts = (torch.sigmoid(predicts) > 0.5).int().detach().cpu()
        targets = data[1].detach().cpu()

        eps = 1e-12

        iou = iou_pytorch(predicts, targets)
        f1 = dice_coeff(predicts, targets)

        self.f1 += f1
        self.iou += iou.sum()
        self.len += 1

    def after_epoch_end(self, val_loss, **ignore):
        self.miou and self._iou()
        self.mf1 and self._f1()

    def _f1(self):
        mf1 = self.f1 / self.len
        png_file = self.scalars(
            {
                'mf1': mf1,
            }, 'mf1'
        )
        if png_file:
            self.update_sheet('f1', {'raw': png_file, 'processor': 'upload_image'})

    def _iou(self):
        miou = torch.sum(self.iou) / 800
        png_file = self.scalars(
            {
                'miou': miou,
            }, 'miou'
        )
        if png_file:
            self.update_sheet('iou', {'raw': png_file, 'processor': 'upload_image'})

df = pd.read_csv('/home/featurize/data/train.csv')
df['image_path'] = '/home/featurize/data/train/img/'
df['mask_path'] = f'/home/featurize/data/train/mask/'

train_df = df[df.fold != CONFIG['fold']].reset_index(drop=True)
val_df = df[df.fold == CONFIG['fold']].reset_index(drop=True)

trainset = AliDataset(train_df, create_transformation('train'),train=True)
valset = AliDataset(val_df, create_transformation('val'))

train_loader = DataLoader(
    trainset,
    batch_size=CONFIG['batch_size'],
    num_workers=8,
    shuffle=True,
    pin_memory=True,
)
val_loader = DataLoader(
    valset,
    batch_size=CONFIG['batch_size'],
    num_workers=8,
    pin_memory=True
)

df.head()

model = smp.UnetPlusPlus(
    CONFIG['model_name'],
    classes=1,
    encoder_weights='imagenet',
    activation=None,
)

optimizer = torch.optim.AdamW(model.parameters(), lr=CONFIG['lr'], betas=(0.9, 0.999), eps=1e-08, amsgrad=False)
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2, eta_min=1e-6, last_epoch=-1)
def after_epoch_end(miner, **payload):
    scheduler.step(miner.current_epoch)

bce_loss = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([3.]).cuda()).cuda()
def forward_fn(trainer, data):
    images, masks = data
    images, masks = images.cuda(), masks.cuda()
    mask = model(images).squeeze(1)
    loss = bce_loss(mask, masks)
    return mask, loss

miner = minetorch.Miner(
    code=os.getenv('CODE', f'fold-{CONFIG["fold"]}'),
    alchemistic_directory=os.getenv('ALCHEMISTIC_DIRECTORY', f'/home/featurize/{CONFIG['DIR']}/{CONFIG["model_name"]}-{CONFIG["img_size"]}'),
    model=model.cuda(),
    forward=forward_fn,
    optimizer=optimizer,
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    max_epochs=CONFIG['epoch'],
    in_notebook=True,
    loss_func=None,
    amp=True,
    plugins=[MultiClassesSegmentationMetricWithLogic()],
    hooks={
        'after_epoch_end': after_epoch_end,
    },
    trival=False,
    resume=False
)

miner.train()

import IPython
IPython.display.Image('/home/featurize/Experiments/timm-efficientnet-b0-224/fold-0/graphs/loss.png')

IPython.display.Image('/home/featurize/Experiments/timm-efficientnet-b0-224/fold-0/graphs/miou.png')

IPython.display.Image('/home/featurize/Experiments/timm-efficientnet-b0-224/fold-0/graphs/mf1.png')