Spaces:
Runtime error
Runtime error
menghanxia
commited on
Commit
•
2a10b61
1
Parent(s):
565e612
fixed the randomness issue
Browse files- app.py +18 -13
- inference.py +3 -7
- models/clusterkit.py +1 -0
app.py
CHANGED
@@ -1,9 +1,12 @@
|
|
1 |
import gradio as gr
|
2 |
import os, requests
|
|
|
3 |
from inference import setup_model, colorize_grayscale, predict_anchors
|
4 |
|
5 |
-
|
6 |
-
|
|
|
|
|
7 |
|
8 |
## step 1: set up model
|
9 |
device = "cpu"
|
@@ -30,7 +33,7 @@ def switch_states(is_checked):
|
|
30 |
demo = gr.Blocks(title="DISCO")
|
31 |
with demo:
|
32 |
gr.Markdown(value="""
|
33 |
-
**Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors
|
34 |
""")
|
35 |
with gr.Row():
|
36 |
with gr.Column(scale=1):
|
@@ -39,13 +42,13 @@ with demo:
|
|
39 |
with gr.Row():
|
40 |
Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)")
|
41 |
Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \
|
42 |
-
label="Colorization resolution", value="Low (256x256)")
|
43 |
Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors')
|
44 |
with gr.Row():
|
45 |
Button_show_anchor = gr.Button(value="Predict anchors", visible=False)
|
46 |
Button_run = gr.Button(value="Colorize")
|
47 |
with gr.Column(scale=1):
|
48 |
-
Image_output = gr.Image(type="numpy", label="Output"
|
49 |
|
50 |
Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor])
|
51 |
Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor)
|
@@ -55,14 +58,16 @@ with demo:
|
|
55 |
gr.Markdown(value="""
|
56 |
**Guideline**
|
57 |
1. upload your image;
|
58 |
-
2.
|
59 |
-
3.
|
60 |
-
-
|
61 |
-
-
|
62 |
""")
|
63 |
gr.HTML(value="""
|
64 |
-
<p style=
|
65 |
""")
|
66 |
-
|
67 |
-
|
68 |
-
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 |
+
RUN_MODE = "remote"
|
7 |
+
if RUN_MODE != "local":
|
8 |
+
os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")
|
9 |
+
os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")
|
10 |
|
11 |
## step 1: set up model
|
12 |
device = "cpu"
|
|
|
33 |
demo = gr.Blocks(title="DISCO")
|
34 |
with demo:
|
35 |
gr.Markdown(value="""
|
36 |
+
**Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors**. Check our project page [*Here*](https://menghanxia.github.io/projects/disco.html).
|
37 |
""")
|
38 |
with gr.Row():
|
39 |
with gr.Column(scale=1):
|
|
|
42 |
with gr.Row():
|
43 |
Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)")
|
44 |
Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \
|
45 |
+
label="Colorization resolution (Low is more stable)", value="Low (256x256)")
|
46 |
Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors')
|
47 |
with gr.Row():
|
48 |
Button_show_anchor = gr.Button(value="Predict anchors", visible=False)
|
49 |
Button_run = gr.Button(value="Colorize")
|
50 |
with gr.Column(scale=1):
|
51 |
+
Image_output = gr.Image(type="numpy", label="Output").style(height=480)
|
52 |
|
53 |
Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor])
|
54 |
Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor)
|
|
|
58 |
gr.Markdown(value="""
|
59 |
**Guideline**
|
60 |
1. upload your image;
|
61 |
+
2. set up the arguments: "Num. of anchors" and "Colorization resolution";
|
62 |
+
3. run the colorization (two modes supported):
|
63 |
+
- *Automatic mode*: click "Colorize" to get the automatically colorized output.
|
64 |
+
- *Editable mode*: check ""Show editable anchors" and click "Predict anchors". Then, modify the colors of the predicted anchors (only anchor region will be used). Finally, click "Colorize" to get the result.
|
65 |
""")
|
66 |
gr.HTML(value="""
|
67 |
+
<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>
|
68 |
""")
|
69 |
+
|
70 |
+
if RUN_MODE == "local":
|
71 |
+
demo.launch(server_name='9.134.253.83',server_port=7788)
|
72 |
+
else:
|
73 |
+
demo.launch()
|
inference.py
CHANGED
@@ -11,17 +11,13 @@ from utils import util
|
|
11 |
|
12 |
|
13 |
def setup_model(checkpt_path, device="cuda"):
|
14 |
-
seed = 130
|
15 |
-
np.random.seed(seed)
|
16 |
-
torch.manual_seed(seed)
|
17 |
-
torch.cuda.manual_seed(seed)
|
18 |
#print('--------------', torch.cuda.is_available())
|
19 |
"""Load the model into memory to make running multiple predictions efficient"""
|
20 |
colorLabeler = basic.ColorLabel(device=device)
|
21 |
colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True, colorLabeler=colorLabeler)
|
22 |
colorizer = colorizer.to(device)
|
23 |
#checkpt_path = "./checkpoints/disco-beta.pth.rar"
|
24 |
-
assert os.path.exists(checkpt_path)
|
25 |
data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
|
26 |
colorizer.load_state_dict(data_dict['state_dict'])
|
27 |
colorizer.eval()
|
@@ -89,8 +85,8 @@ def predict_anchors(colorizer, color_class, rgb_img, n_anchors, is_high_res, is_
|
|
89 |
n_anchors = min(n_anchors, 14)
|
90 |
target_res = (512,512) if is_high_res else (256,256)
|
91 |
input_grays, input_colors, org_grays = prepare_data(rgb_img, target_res)
|
92 |
-
input_grays = input_grays.
|
93 |
-
input_colors = input_colors.
|
94 |
|
95 |
sampled_T, sp_size = 0, 16
|
96 |
pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \
|
|
|
11 |
|
12 |
|
13 |
def setup_model(checkpt_path, device="cuda"):
|
|
|
|
|
|
|
|
|
14 |
#print('--------------', torch.cuda.is_available())
|
15 |
"""Load the model into memory to make running multiple predictions efficient"""
|
16 |
colorLabeler = basic.ColorLabel(device=device)
|
17 |
colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True, colorLabeler=colorLabeler)
|
18 |
colorizer = colorizer.to(device)
|
19 |
#checkpt_path = "./checkpoints/disco-beta.pth.rar"
|
20 |
+
assert os.path.exists(checkpt_path), "No checkpoint found!"
|
21 |
data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
|
22 |
colorizer.load_state_dict(data_dict['state_dict'])
|
23 |
colorizer.eval()
|
|
|
85 |
n_anchors = min(n_anchors, 14)
|
86 |
target_res = (512,512) if is_high_res else (256,256)
|
87 |
input_grays, input_colors, org_grays = prepare_data(rgb_img, target_res)
|
88 |
+
input_grays = input_grays.to(device)
|
89 |
+
input_colors = input_colors.to(device)
|
90 |
|
91 |
sampled_T, sp_size = 0, 16
|
92 |
pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \
|
models/clusterkit.py
CHANGED
@@ -103,6 +103,7 @@ def initialize(X, num_clusters):
|
|
103 |
:param num_clusters: (int) number of clusters
|
104 |
:return: (np.array) initial state
|
105 |
"""
|
|
|
106 |
num_samples = len(X)
|
107 |
indices = np.random.choice(num_samples, num_clusters, replace=False)
|
108 |
initial_state = X[indices]
|
|
|
103 |
:param num_clusters: (int) number of clusters
|
104 |
:return: (np.array) initial state
|
105 |
"""
|
106 |
+
np.random.seed(1)
|
107 |
num_samples = len(X)
|
108 |
indices = np.random.choice(num_samples, num_clusters, replace=False)
|
109 |
initial_state = X[indices]
|