File size: 875 Bytes
fac6837
 
 
 
 
 
fb6a1c2
fac6837
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
96b0185
 
 
 
 
 
 
fd7ddbc
 
 
fac6837
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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):
    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",
            default_selected="eccv16",
            label="model"
        )
    ],
    [
        gr.outputs.Image(label="Output")
    ]
)