import os
import cv2
import numpy as np
import torch
import gradio as gr
import torchvision.transforms as transforms

os.system("pip freeze")

model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', weights='DEFAULT')
model.eval()


def image_to_tensor(image):
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])(image)


def make_transparent_foreground(image, mask):
    # split the image into channels
    b, g, r = cv2.split(np.array(image).astype('uint8'))
    # add an alpha channel with and fill all with transparent pixels (max 255)
    a = np.ones(mask.shape, dtype='uint8') * 255
    # merge the alpha channel back
    alpha_im = cv2.merge([b, g, r, a], 4)
    # create a transparent background
    bg = np.zeros(alpha_im.shape)
    # set up the new mask
    new_mask = np.stack([mask, mask, mask, mask], axis=2)
    # copy only the foreground color pixels from the original image where mask is set
    return np.where(new_mask, alpha_im, bg).astype(np.uint8)


def makeMask(image):
    input_tensor = image_to_tensor(image)
    input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model

    # move the input and model to GPU for speed if available
    if torch.cuda.is_available():
        input_batch = input_batch.to('cuda')
        model.to('cuda')

    with torch.no_grad():
        output = model(input_batch)['out'][0]
    output_predictions = output.argmax(0)

    # create a binary (black and white) mask of the profile foreground
    mask = output_predictions.byte().cpu().numpy()
    background = np.zeros(mask.shape)
    return np.where(mask, 255, background).astype(np.uint8)


def predict(image):
    mask = makeMask(image)
    return make_transparent_foreground(image, mask)


title = "Zero Background"
description = r"""
## Remove image background

This is another implementation of <a href='https://github.com/eugenesiow/practical-ml/blob/master/notebooks/Remove_Image_Background_DeepLabV3.ipynb' target='_blank'>eugenesiow</a>.
It has no any particular purpose than start research on AI models.

"""

article = r"""
Questions, doubts, comments, please email 📧 `leonelhs@gmail.com`

This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/leonelhs/face-shine' target='_blank'>Github ⭐</a>

<a href="https://www.buymeacoffee.com/leonelhs">
<img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" />
</a>

<center><img src='https://visitor-badge.glitch.me/badge?page_id=deoldify.visitor-badge' alt='visitor badge'></center>
"""

demo = gr.Interface(
    predict, [
        gr.Image(type="pil", label="Image"),
    ], [
        gr.Image(type="pil", label="Image alpha background")
    ],
    title=title,
    description=description,
    article=article)

demo.queue().launch()