import gradio as gr
from PIL import Image, PngImagePlugin
import numpy as np
import os

os.system("chmod +x models/waifu2x-ncnn-vulkan")

noisedict = {
    "なし": -1,
    "低": 0,
    "中": 1,
    "高": 2,
    "最高": 3
}

scaledict = {
    "1倍": 1,
    "2倍": 2,
}

formatdict = {
    "PNG": "png",
    "JPG": "jpg",
    "WebP": "webp"
}

def response_greet(image, noise, scale, format):
    info = image.info
    n = noisedict[noise]
    s = scaledict[scale]
    f = formatdict[format]
    image.save("input.png")
    os.system(f"models/waifu2x-ncnn-vulkan -i input.png -o output.{f} -n {n} -s {s} -f {f} -g -1")
    image = Image.open(f"output.{f}")
    image.info = info
    return image
    
with gr.Blocks() as app:
    gr.Markdown("## Waifu2x with png metadata demo")
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="入力画像", interactive=True, type="pil", )
            noise = gr.Radio(choices=["なし", "低", "中", "高", "最高"], label="ノイズ除去", value="中", interactive=True, type="value"),
            scale = gr.Radio(choices=["1倍", "2倍"], label="拡大", value="2倍", interactive=True, type="value"),
            format = gr.Radio(choices=["PNG", "JPG", "WebP"], label="出力フォーマット（※現時点ではPNGのみ選択できます）", value="PNG", type="value"),
            button = gr.Button("送信")
        with gr.Column():
            output = gr.Image(label="出力画像", type="pil")
    button.click(fn=response_greet, inputs=[image, noise[0], scale[0], format[0]], outputs=output, api_name="Waifu2xで画像をアップコンバートします。")

app.launch()