import requests
import gradio as gr

import paddle
from paddleseg.cvlibs import Config

from matting.core import predict
from matting.model import *
from matting.dataset import MattingDataset


def download_file(http_address, file_name):
    r = requests.get(http_address, allow_redirects=True)
    open(file_name, 'wb').write(r.content)

cfg_paths = ['configs/modnet/modnet_mobilenetv2.yml', 'configs/modnet/modnet_resnet50_vd.yml', 'configs/modnet/modnet_hrnet_w18.yml']
cfgs = [Config(cfg) for cfg in cfg_paths]

download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-mobilenetv2.pdparams', 'modnet-mobilenetv2.pdparams')
download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-resnet50_vd.pdparams', 'modnet-resnet50_vd.pdparams')
download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-hrnet_w18.pdparams', 'modnet-hrnet_w18.pdparams')
models_paths = ['modnet-mobilenetv2.pdparams', 'modnet-resnet50_vd.pdparams', 'modnet-hrnet_w18.pdparams']
models = [cfg.model for cfg in cfgs]


def inference(image, chosen_model):
    paddle.set_device('cpu')

    cfg = cfgs[chosen_model]
    val_dataset = cfg.val_dataset
    img_transforms = val_dataset.transforms

    model = models[chosen_model]

    alpha_pred = predict(model,
                         model_path=models_paths[chosen_model],
                         transforms=img_transforms,
                         image_list=[image])

    return alpha_pred


inputs = [gr.inputs.Image(label='Input Image'),
          gr.inputs.Radio(['MobileNetV2', 'ResNet50_vd', 'HRNet_W18'], label='Model', type='index')]

gr.Interface(
    inference, 
    inputs,
    gr.outputs.Image(label='Output'),
    title='PaddleSeg - Matting',
    examples=[['images/armchair.jpg', 'MobileNetV2'],
              ['images/cat.jpg', 'ResNet50_vd'],
              ['images/plant.jpg', 'HRNet_W18']]
    ).launch()