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 = "siggraph17"):
    if model == "eccv16":
        img = siggraph17(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

css='''
.Box {
	background-color: var(--color-canvas-default);
	border-color: var(--color-border-default);
	border-style: solid;
	border-width: 1px;
	border-radius: 6px;
}
.d-flex {
	display: flex !important;
}
.flex-md-row {
	flex-direction: row !important;
}
.flex-column {
	flex-direction: column !important;
}
'''
title = "Image Colorization Using AI Models"
description = r"""<center>An automatic colorization functionality for Real-Time User-Guided Image Colorization with Learned Deep Priors,ECCV16 & SIGGRAPH 2017 Models!<br>
Practically the algorithm is used to COLORIZE your **old BLACK & WHITE / GRAYSCALE photos**.<br>
To use it, simply just upload the concerned image.<br>
"""
article = r"""


"""

#with gr.Interface(css=css) as mainBody:
gr.HTML("""<style>""" + css+ """</Style>""")
  
mainBody = gr.Interface(
    interface,
    [
        gr.components.Image(type="pil", label="image"),
        gr.components.Radio(
            ["siggraph17"],
            type="value",
            label="model"
        )
    ],
    [
        gr.components.Image(label="output")
    ],
    #inputs="sketchpad",
    #outputs="label",
    theme="huggingface",
    title=title,
    description=description,
    article=article,
    live=True,
)
mainBody.launch()