Spaces:
Running
Running
| import gradio as gr | |
| import os, requests | |
| from inference import setup_model, colorize_grayscale, predict_anchors | |
| ## download checkpoint | |
| def download_file_from_google_drive(id, destination): | |
| def get_confirm_token(response): | |
| for key, value in response.cookies.items(): | |
| if key.startswith('download_warning'): | |
| return value | |
| return None | |
| def save_response_content(response, destination): | |
| CHUNK_SIZE = 32768 | |
| with open(destination, "wb") as f: | |
| for chunk in response.iter_content(CHUNK_SIZE): | |
| if chunk: # filter out keep-alive new chunks | |
| f.write(chunk) | |
| URL = "https://docs.google.com/uc?export=download" | |
| session = requests.Session() | |
| response = session.get(URL, params = { 'id' : id }, stream = True) | |
| token = get_confirm_token(response) | |
| if token: | |
| params = { 'id' : id, 'confirm' : token } | |
| response = session.get(URL, params = params, stream = True) | |
| save_response_content(response, destination) | |
| id = "1J4vB6kG4xBLUUKpXr5IhnSSa4maXgRvQ" | |
| destination = "disco-beta.pth.rar" | |
| download_file_from_google_drive(id, destination) | |
| os.rename("disco-beta.pth.tar", "./checkpoints/disco-beta.pth.tar") | |
| ## step 1: set up model | |
| device = "cuda" | |
| checkpt_path = "./checkpoints/disco-beta.pth.tar" | |
| assert os.path.exists(checkpt_path), "No checkpoint found!" | |
| colorizer, colorLabeler = setup_model(checkpt_path, device=device) | |
| def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable): | |
| if hint_img is None: | |
| hint_img = rgb_img | |
| output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, is_high_res, is_editable, device) | |
| return output | |
| def click_predanchors(rgb_img, n_anchors, is_high_res, is_editable): | |
| output = predict_anchors(colorizer, colorLabeler, rgb_img, n_anchors, is_high_res, is_editable, device) | |
| return output | |
| ## step 2: configure interface | |
| def switch_states(is_checked): | |
| if is_checked: | |
| return gr.Image.update(visible=True), gr.Button.update(visible=True) | |
| else: | |
| return gr.Image.update(visible=False), gr.Button.update(visible=False) | |
| demo = gr.Blocks(title="DISCO: Image Colorization") | |
| with demo: | |
| gr.Markdown(value="""**DISCO: image colorization that disentangles color multimodality and spatial affinity via global anchors**.""") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| Image_input = gr.Image(type="numpy", label="Input", interactive=True) | |
| Image_anchor = gr.Image(type="numpy", label="Anchor", tool="color-sketch", interactive=True, visible=False) | |
| with gr.Row(): | |
| Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)") | |
| Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \ | |
| label="Colorization resolution", value="Low (256x256)") | |
| Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors') | |
| with gr.Row(): | |
| Button_show_anchor = gr.Button(value="Predict anchors", visible=False) | |
| Button_run = gr.Button(value="Colorize") | |
| with gr.Column(scale=1): | |
| Image_output = gr.Image(type="numpy", label="Output", shape=[100,100]) | |
| Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor]) | |
| Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor) | |
| Button_run.click(fn=click_colorize, inputs=[Image_input, Image_anchor, Num_anchor, Radio_resolution, Ckeckbox_editable], \ | |
| outputs=Image_output) | |
| ## guiline | |
| gr.Markdown(value=""" | |
| **Guideline** | |
| 1. Upload your image; | |
| 2. Set up the arguments: "Num. of anchors" and "Colorization resolution"; | |
| 3. Two modes are supported: | |
| - **Editable**: check ""Show editable anchors" and click "Predict anchors". Then, modify the colors of the predicted anchors (anchor mask will be applied afterward). Finally, click "Colorize" to get the result. | |
| - **Automatic**: click "Colorize" to get the automatically colorized output. | |
| *To know more about the method, please refer to our project page: [https://menghanxia.github.io/projects/disco.html](https://menghanxia.github.io/projects/disco.html)* | |
| """) | |
| demo.launch(server_name='9.134.253.83',server_port=7788) |