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