import os
import gradio as gr
from PIL import Image

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

def inference(img, model):
    os.system('mkdir test')
    #basewidth = 256
    #wpercent = (basewidth / float(img.size[0]))
    #hsize = int((float(img.size[1]) * float(wpercent)))
    #img = img.resize((basewidth, hsize), Image.ANTIALIAS)
    width, height = img.size
    newsize= (width, height)
    img.save("test/1.png", "PNG")
    if model == 'LOL':
        os.system('python main_test_HWMNet.py --input_dir test --weights experiments/pretrained_models/LOL_enhancement_HWMNet.pth')
    elif model == 'MIT-5K':
        os.system('python main_test_HWMNet.py --input_dir test --weights experiments/pretrained_models/MIT5K_enhancement_HWMNet.pth')
    
    im = Image.open('result/1.png')
    im = im.resize(newsize)
    return 'result/1.png'


title = "Low-light Image Enhancement"
description = "Demo di uno script python che usa pythorch, opencv e numpy per eseguire il \"low light image enhancemnent\". Questa tecnica consente di migliorare graficamente un immagine con poca luce avvalendosi di dataset come \"LOL\" e \"MIT-Adobe FiveK\". <br> \"LOL\" è più adatto per immagini molto scure, di conseguenza è più invasivo. \"MIT-Adobe FiveK\" è meno invasivo e quindi adatto ad immagini più luminose di base."
article = "<p style='text-align: center'>made by Peter Minerba for high-scool final exam ツ </p>"
examples = [['low-light.png', 'LOL'], ['low-light_2.png', 'MIT-5K']]
gr.Interface(
    inference,
    [gr.inputs.Image(type="pil", label="Input"), gr.inputs.Dropdown(choices=['LOL', 'MIT-5K'], type="value", default='LOL', label="model")],
    gr.outputs.Image(type="file", label="Output"),
    title=title,
    description=description,
    article=article,
    allow_flagging=False,
    allow_screenshot=False,
    examples=examples,
    css='#2 {max-width: 100%; max-height: 70%;}'
).launch(debug=True)