import gradio as gr
import numpy as np
import colorizers as c

from colorizers.util import postprocess_tens, preprocess_img

def interface(image, model: str = "eccv16"):
    if model == "eccv16":
        img = c.eccv16(pretrained=True).eval()
    else:
        img = c.siggraph17(pretrained=True).eval()
    oimg = np.asarray(image)
    if(oimg.ndim == 2):
        oimg = np.tile(oimg[:,:,None], 3)
    (tens_l_orig, tens_l_rs) = preprocess_img(oimg)

    output_img = postprocess_tens(
        tens_l_orig,
        img(tens_l_rs).cpu()
    )
    return output_img

gr.Interface(
    interface,
    [
        gr.components.Image(type="pil", label="image"),
        gr.components.Radio(
            ["eccv16", "siggraph17"],
            type="value",
            label="model"
        )
    ],
    [
        gr.components.Image(label="output")
    ]
).launch()