Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,6 @@ import gradio as gr
|
|
3 |
from PIL import Image
|
4 |
import torch
|
5 |
import matplotlib.pyplot as plt
|
6 |
-
import cv2
|
7 |
import torch
|
8 |
import numpy as np
|
9 |
|
@@ -11,7 +10,7 @@ processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
|
|
11 |
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
|
12 |
|
13 |
|
14 |
-
def process_image(image, prompt
|
15 |
inputs = processor(
|
16 |
text=prompt, images=image, padding="max_length", return_tensors="pt"
|
17 |
)
|
@@ -32,35 +31,33 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
|
|
32 |
mask_min = mask.min()
|
33 |
mask_max = mask.max()
|
34 |
mask = (mask - mask_min) / (mask_max - mask_min)
|
|
|
35 |
|
36 |
-
# threshold the mask
|
37 |
-
bmask = mask > threhsold
|
38 |
-
# zero out values below the threshold
|
39 |
-
mask[mask < threhsold] = 0
|
40 |
-
|
41 |
-
fig, ax = plt.subplots()
|
42 |
-
ax.imshow(image)
|
43 |
-
ax.imshow(mask, alpha=alpha_value, cmap="jet")
|
44 |
-
|
45 |
-
if draw_rectangles:
|
46 |
-
contours, hierarchy = cv2.findContours(
|
47 |
-
bmask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
|
48 |
-
)
|
49 |
-
for contour in contours:
|
50 |
-
x, y, w, h = cv2.boundingRect(contour)
|
51 |
-
rect = plt.Rectangle(
|
52 |
-
(x, y), w, h, fill=False, edgecolor="yellow", linewidth=2
|
53 |
-
)
|
54 |
-
ax.add_patch(rect)
|
55 |
|
56 |
-
|
57 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
|
63 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
64 |
|
65 |
|
66 |
title = "Interactive demo: zero-shot image segmentation with CLIPSeg"
|
@@ -72,51 +69,36 @@ with gr.Blocks() as demo:
|
|
72 |
gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
|
73 |
gr.Markdown(article)
|
74 |
gr.Markdown(description)
|
75 |
-
gr.Markdown(
|
76 |
-
"*Example images are taken from the [ImageNet-A](https://paperswithcode.com/dataset/imagenet-a) dataset*"
|
77 |
-
)
|
78 |
|
79 |
with gr.Row():
|
80 |
with gr.Column():
|
81 |
input_image = gr.Image(type="pil")
|
82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
83 |
input_slider_T = gr.Slider(
|
84 |
minimum=0, maximum=1, value=0.4, label="Threshold"
|
85 |
)
|
86 |
-
input_slider_A = gr.Slider(minimum=0, maximum=1, value=0.5, label="Alpha")
|
87 |
-
draw_rectangles = gr.Checkbox(label="Draw rectangles")
|
88 |
btn_process = gr.Button(label="Process")
|
89 |
|
90 |
with gr.Column():
|
91 |
-
|
92 |
output_mask = gr.Image(label="Mask")
|
93 |
-
output_image = gr.Image(label="Output Image")
|
94 |
|
95 |
btn_process.click(
|
96 |
-
|
97 |
inputs=[
|
|
|
|
|
98 |
input_image,
|
99 |
-
input_prompt,
|
100 |
input_slider_T,
|
101 |
-
input_slider_A,
|
102 |
-
draw_rectangles,
|
103 |
],
|
104 |
-
outputs=[
|
105 |
)
|
106 |
|
107 |
-
gr.Examples(
|
108 |
-
[
|
109 |
-
["0.003473_cliff _ cliff_0.51112.jpg", "dog", 0.5, 0.5, True],
|
110 |
-
["0.001861_submarine _ submarine_0.9862991.jpg", "beacon", 0.55, 0.4, True],
|
111 |
-
["0.004658_spatula _ spatula_0.35416836.jpg", "banana", 0.4, 0.5, True],
|
112 |
-
],
|
113 |
-
inputs=[
|
114 |
-
input_image,
|
115 |
-
input_prompt,
|
116 |
-
input_slider_T,
|
117 |
-
input_slider_A,
|
118 |
-
draw_rectangles,
|
119 |
-
],
|
120 |
-
)
|
121 |
|
122 |
-
demo.launch()
|
|
|
3 |
from PIL import Image
|
4 |
import torch
|
5 |
import matplotlib.pyplot as plt
|
|
|
6 |
import torch
|
7 |
import numpy as np
|
8 |
|
|
|
10 |
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
|
11 |
|
12 |
|
13 |
+
def process_image(image, prompt):
|
14 |
inputs = processor(
|
15 |
text=prompt, images=image, padding="max_length", return_tensors="pt"
|
16 |
)
|
|
|
31 |
mask_min = mask.min()
|
32 |
mask_max = mask.max()
|
33 |
mask = (mask - mask_min) / (mask_max - mask_min)
|
34 |
+
return mask
|
35 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
36 |
|
37 |
+
def get_masks(prompts, img, threhsold):
|
38 |
+
prompts = prompts.split(",")
|
39 |
+
masks = []
|
40 |
+
for prompt in prompts:
|
41 |
+
mask = process_image(img, prompt)
|
42 |
+
mask = mask > threhsold
|
43 |
+
masks.append(mask)
|
44 |
+
return masks
|
45 |
+
|
46 |
|
47 |
+
def extract_image(pos_prompts, neg_prompts, img, threhsold):
|
48 |
+
positive_masks = get_masks(pos_prompts, img, 0.5)
|
49 |
+
negative_masks = get_masks(neg_prompts, img, 0.5)
|
50 |
|
51 |
+
# combine masks into one masks, logic OR
|
52 |
+
pos_mask = np.any(np.stack(positive_masks), axis=0)
|
53 |
+
neg_mask = np.any(np.stack(negative_masks), axis=0)
|
54 |
+
final_mask = pos_mask & ~neg_mask
|
55 |
+
|
56 |
+
# extract the final image
|
57 |
+
final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255, "L")
|
58 |
+
output_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
|
59 |
+
output_image.paste(img, mask=final_mask)
|
60 |
+
return output_image, final_mask
|
61 |
|
62 |
|
63 |
title = "Interactive demo: zero-shot image segmentation with CLIPSeg"
|
|
|
69 |
gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
|
70 |
gr.Markdown(article)
|
71 |
gr.Markdown(description)
|
|
|
|
|
|
|
72 |
|
73 |
with gr.Row():
|
74 |
with gr.Column():
|
75 |
input_image = gr.Image(type="pil")
|
76 |
+
positive_prompts = gr.Textbox(
|
77 |
+
label="Please describe what you want to identify (comma separated)"
|
78 |
+
)
|
79 |
+
negative_prompts = gr.Textbox(
|
80 |
+
label="Please describe what you want to ignore (comma separated)"
|
81 |
+
)
|
82 |
+
|
83 |
input_slider_T = gr.Slider(
|
84 |
minimum=0, maximum=1, value=0.4, label="Threshold"
|
85 |
)
|
|
|
|
|
86 |
btn_process = gr.Button(label="Process")
|
87 |
|
88 |
with gr.Column():
|
89 |
+
output_image = gr.Image(label="Result")
|
90 |
output_mask = gr.Image(label="Mask")
|
|
|
91 |
|
92 |
btn_process.click(
|
93 |
+
extract_image,
|
94 |
inputs=[
|
95 |
+
positive_prompts,
|
96 |
+
negative_prompts,
|
97 |
input_image,
|
|
|
98 |
input_slider_T,
|
|
|
|
|
99 |
],
|
100 |
+
outputs=[output_image, output_mask],
|
101 |
)
|
102 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
103 |
|
104 |
+
demo.launch()
|