menghanxia/disco

1import gradio as gr2import os, requests3import numpy as np4from inference import setup_model, colorize_grayscale, predict_anchors5 6## local |  remote7RUN_MODE = "remote"8if RUN_MODE != "local":9    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")10    os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")11    ## examples12    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/01.jpg")13    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/02.jpg")14    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/03.jpg")15    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/04.jpg")16 17## step 1: set up model18device = "cpu"19checkpt_path = "checkpoints/disco-beta.pth.rar"20colorizer, colorLabeler = setup_model(checkpt_path, device=device)21 22def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable):23    if hint_img is None:24        hint_img = rgb_img25    output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, is_high_res, is_editable, device)26    return output27 28def click_predanchors(rgb_img, n_anchors, is_high_res, is_editable):29    output = predict_anchors(colorizer, colorLabeler, rgb_img, n_anchors, is_high_res, is_editable, device)30    return output31 32## step 2: configure interface33def switch_states(is_checked):34    if is_checked:35        return gr.Image.update(visible=True), gr.Button.update(visible=True)36    else:37        return gr.Image.update(visible=False), gr.Button.update(visible=False)38 39demo = gr.Blocks(title="DISCO")40with demo:41    gr.Markdown(value="""42                    **Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors**. Check our [project page](https://menghanxia.github.io/projects/disco.html) 😛.43                    """)44    with gr.Row():45        with gr.Column():46            with gr.Row().style(height=480):47                Image_input = gr.Image(type="numpy", label="Input", interactive=True)48                Image_anchor = gr.Image(type="numpy", label="Anchor", tool="color-sketch", interactive=True, visible=False)49            with gr.Row():50                Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)")51                Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \52                                                label="Colorization resolution (Low is more stable)", value="Low (256x256)")53            with gr.Row():54                Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors')55                Button_show_anchor = gr.Button(value="Predict anchors", visible=False)56            Button_run = gr.Button(value="Colorize")57        with gr.Column():58            Image_output = gr.Image(type="numpy", label="Output").style(height=480)59 60    Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor])61    Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor)62    Button_run.click(fn=click_colorize, inputs=[Image_input, Image_anchor, Num_anchor, Radio_resolution, Ckeckbox_editable], \63                    outputs=Image_output)64    65    ## guiline66    gr.Markdown(value="""    67                    🔔**Guideline**68                    1. Upload your image or select one from the examples.69                    2. Set up the arguments: "Num. of anchors" and "Colorization resolution".70                    3. Run the colorization (two modes supported):71                        - 📀Automatic mode: **Click** "Colorize" to get the automatically colorized output.72                        - ✏️Editable mode: **Check** ""Show editable anchors"; **Click** "Predict anchors"; **Redraw** the anchor colors (only anchor region will be used); **Click** "Colorize" to get the result.73                    """)74    if RUN_MODE != "local":75        gr.Examples(examples=[76                    ['01.jpg', 8, "Low (256x256)"],77                    ['02.jpg', 8, "Low (256x256)"],78                    ['03.jpg', 8, "Low (256x256)"],79                    ['04.jpg', 8, "Low (256x256)"],80                    ], 81                    inputs=[Image_input,Num_anchor,Radio_resolution], outputs=[Image_output], label="Examples")82    gr.HTML(value="""83                <p style="text-align:center; color:orange"><a href='https://menghanxia.github.io/projects/disco.html' target='_blank'>DISCO Project Page</a> | <a href='https://github.com/MenghanXia/DisentangledColorization' target='_blank'>Github Repo</a></p>84                    """)85 86if RUN_MODE == "local":87    demo.launch(server_name='9.134.253.83',server_port=7788)88else:89    demo.launch()