import gradio as gr
import numpy as np
import colorizers as c
from PIL import Image

from colorizers.util import postprocess_tens, preprocess_img

def interface(image: Image, model: str):
    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.inputs.Image(type="pil", label="Image"),
        gr.inputs.Radio([
            "eccv16", "siggraph17"
        ], type="value", default="eccv16", label="Model")
    ]
)