from PIL import Image
from scipy import misc, ndimage
import matplotlib.pyplot as plt
import torch
import numpy as np
from torchvision import models
from torchvision import transforms


def intepret_semantic_model(model, device, alpha=50):
    invTrans = transforms.Compose([
        transforms.Normalize(mean=[0., 0., 0.],
                             std=[1 / 0.229, 1 / 0.224, 1 / 0.225]),
        transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1., 1., 1.]),
    ])
    preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((512, 512)),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    model = model.to(device).eval()
    model.load()
    image = Image.open("../etc/dash.jpg")
    preproc_img = preprocess(image)

    preproc_img = preproc_img.unsqueeze(0).to(device)

    with torch.no_grad():
        out = model(preproc_img)
        mask = out[0].permute(1, 2, 0)
        mask = mask.mul(255).clamp(0, 255)
        mask = mask.detach().cpu().numpy().astype(np.float32)

    preproc_img = invTrans(preproc_img)
    plt.figure()
    plt.axis('off')
    plt.imshow(preproc_img[0].permute(1, 2, 0).cpu())
    plt.imshow(apply_sharpen_filter(mask, alpha),
               alpha=0.4,
               cmap='winter',
               interpolation='gaussian')
    plt.show()


def apply_sharpen_filter(img, alpha):
    blurred_filter = ndimage.gaussian_filter(img, 3)
    filter_blurred = ndimage.gaussian_filter(blurred_filter, 1)
    img = blurred_filter + alpha * (blurred_filter - filter_blurred)
    return img

from PIL import Image
from scipy import misc, ndimage
import matplotlib.pyplot as plt
import torch
import numpy as np
from torchvision import models
from torchvision import transforms


def intepret_semantic_model(model, device, alpha=50):
    invTrans = transforms.Compose([
        transforms.Normalize(mean=[0., 0., 0.],
                             std=[1 / 0.229, 1 / 0.224, 1 / 0.225]),
        transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1., 1., 1.]),
    ])
    preprocess = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((512, 512)),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    model = model.to(device).eval()
    model.load()
    image = Image.open("../etc/dash.jpg")
    preproc_img = preprocess(image)

    preproc_img = preproc_img.unsqueeze(0).to(device)

    with torch.no_grad():
        out = model(preproc_img)
        mask = out[0].permute(1, 2, 0)
        mask = mask.mul(255).clamp(0, 255)
        mask = mask.detach().cpu().numpy().astype(np.float32)

    preproc_img = invTrans(preproc_img)
    plt.figure()
    plt.axis('off')
    plt.imshow(preproc_img[0].permute(1, 2, 0).cpu())
    plt.imshow(apply_sharpen_filter(mask, alpha),
               alpha=0.4,
               cmap='winter',
               interpolation='gaussian')
    plt.show()


def apply_sharpen_filter(img, alpha):
    blurred_filter = ndimage.gaussian_filter(img, 3)
    filter_blurred = ndimage.gaussian_filter(blurred_filter, 1)
    img = blurred_filter + alpha * (blurred_filter - filter_blurred)
    return img
