|
|
|
|
|
from __future__ import annotations |
|
|
|
import json |
|
import shlex |
|
import subprocess |
|
|
|
import gradio as gr |
|
|
|
|
|
def run(image_path: str, class_index: int, sigma_y: float) -> str: |
|
out_name = image_path.split("/")[-1].split(".")[0] |
|
subprocess.run( |
|
shlex.split( |
|
f"python main.py --config confs/inet256.yml --deg colorization --scale 1 --class {class_index} --path_y {image_path} --save_path {out_name} --sigma_y {sigma_y}" |
|
), |
|
cwd="DDNM/hq_demo", |
|
) |
|
return f"DDNM/hq_demo/results/{out_name}/final/00000.png" |
|
|
|
|
|
def create_demo(): |
|
examples = [ |
|
[ |
|
"sample_images/monarch_gray.png", |
|
"monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", |
|
0, |
|
], |
|
[ |
|
"sample_images/tiger_gray.png", |
|
"tiger, Panthera tigris", |
|
0, |
|
], |
|
] |
|
|
|
with open("imagenet_classes.json") as f: |
|
imagenet_class_names = json.load(f) |
|
|
|
with gr.Blocks() as demo: |
|
with gr.Row(): |
|
with gr.Column(): |
|
image = gr.Image(label="Input image", type="filepath") |
|
class_index = gr.Dropdown(label="Class name", choices=imagenet_class_names, type="index", value=950) |
|
sigma_y = gr.Number(label="sigma_y", value=0, precision=2) |
|
run_button = gr.Button("Run") |
|
with gr.Column(): |
|
result = gr.Image(label="Result", type="filepath") |
|
|
|
gr.Examples( |
|
examples=examples, |
|
inputs=[ |
|
image, |
|
class_index, |
|
sigma_y, |
|
], |
|
) |
|
|
|
run_button.click( |
|
fn=run, |
|
inputs=[ |
|
image, |
|
class_index, |
|
sigma_y, |
|
], |
|
outputs=result, |
|
) |
|
return demo |
|
|
