menghanxia commited on
Commit
2a08d14
1 Parent(s): 7d153b0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +88 -88
app.py CHANGED
@@ -1,89 +1,89 @@
1
- import gradio as gr
2
- import os, requests
3
- import numpy as np
4
- from inference import setup_model, colorize_grayscale, predict_anchors
5
-
6
- ## local | remote
7
- RUN_MODE = "remote"
8
- if 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
- ## examples
12
- 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 model
18
- device = "cpu"
19
- checkpt_path = "checkpoints/disco-beta.pth.rar"
20
- colorizer, colorLabeler = setup_model(checkpt_path, device=device)
21
-
22
- def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable):
23
- if hint_img is None:
24
- hint_img = rgb_img
25
- output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, is_high_res, is_editable, device)
26
- return output
27
-
28
- def 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 output
31
-
32
- ## step 2: configure interface
33
- def 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
-
39
- demo = gr.Blocks(title="DISCO")
40
- with 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():
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
- ## guiline
66
- 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
-
86
- if RUN_MODE == "local":
87
- demo.launch(server_name='9.134.253.83',server_port=7788)
88
- else:
89
  demo.launch()
 
1
+ import gradio as gr
2
+ import os, requests
3
+ import numpy as np
4
+ from inference import setup_model, colorize_grayscale, predict_anchors
5
+
6
+ ## local | remote
7
+ RUN_MODE = "remote"
8
+ if 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
+ ## examples
12
+ 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 model
18
+ device = "cpu"
19
+ checkpt_path = "checkpoints/disco-beta.pth.rar"
20
+ colorizer, colorLabeler = setup_model(checkpt_path, device=device)
21
+
22
+ def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable):
23
+ if hint_img is None:
24
+ hint_img = rgb_img
25
+ output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, is_high_res, is_editable, device)
26
+ return output
27
+
28
+ def 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 output
31
+
32
+ ## step 2: configure interface
33
+ def 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
+
39
+ demo = gr.Blocks(title="DISCO")
40
+ with 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():
47
+ Image_input = gr.Image(type="numpy", label="Input", interactive=True).style(height=480)
48
+ Image_anchor = gr.Image(type="numpy", label="Anchor", tool="color-sketch", interactive=True, visible=False).style(height=480)
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
+ ## guiline
66
+ 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
+
86
+ if RUN_MODE == "local":
87
+ demo.launch(server_name='9.134.253.83',server_port=7788)
88
+ else:
89
  demo.launch()