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()