import os
import gradio as gr
from PIL import Image
import torch

os.system(
    'wget https://github.com/FanChiMao/CMFNet/releases/download/v0.0/dehaze_I_OHaze_CMFNet.pth -P experiments/pretrained_models')


def inference(img):
    if not os.path.exists('test'):
        os.system('mkdir test')
    
    basewidth = 512
    wpercent = (basewidth / float(img.size[0]))
    hsize = int((float(img.size[1]) * float(wpercent)))
    img = img.resize((basewidth, hsize), Image.BILINEAR)
    img.save("test/1.png", "PNG")
    os.system(
        'python main_test_CMFNet.py --input_dir test --weights experiments/pretrained_models/dehaze_I_OHaze_CMFNet.pth')
    return 'results/1.png'


title = "Compound Multi-branch Feature Fusion for Image Restoration (Dehaze)"
description = "Gradio demo for CMFNet. CMFNet achieves competitive performance on three tasks: image deblurring, image dehazing and image deraindrop. Here, we provide a demo for image dehaze. To use it, simply upload your image, or click one of the examples to load them. Reference from: https://huggingface.co/akhaliq"
article = "<p style='text-align: center'><a href='https://' target='_blank'>Compound Multi-branch Feature Fusion for Real Image Restoration</a> | <a href='https://github.com/FanChiMao/CMFNet' target='_blank'>Github Repo</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=52Hz_CMFNet_dehazing' alt='visitor badge'></center>"

examples = [['Haze.png']]
gr.Interface(
    inference,
    [gr.components.Image(type="pil", label="Input")],
    gr.components.Image(type="filepath", label="Output"),
    title=title,
    description=description,
    examples=examples
).launch(debug=True)