mkthoma commited on
Commit
9ba6a01
1 Parent(s): 65a8194

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +222 -0
app.py ADDED
@@ -0,0 +1,222 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ultralytics import YOLO
2
+ import gradio as gr
3
+ import torch
4
+ from utils.tools_gradio import fast_process
5
+ from utils.tools import format_results, box_prompt, point_prompt, text_prompt
6
+ from PIL import ImageDraw
7
+ import numpy as np
8
+
9
+ # Load the pre-trained model
10
+ model = YOLO('./weights/FastSAM.pt')
11
+
12
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
+
14
+ # Description
15
+ title = "<center><strong><font size='10'> Fast Segment Anything </font></strong></center>"
16
+
17
+
18
+ examples = [["examples/sa_8776.jpg"], ["examples/sa_414.jpg"], ["examples/sa_1309.jpg"], ["examples/sa_11025.jpg"],
19
+ ["examples/sa_561.jpg"], ["examples/sa_192.jpg"], ["examples/sa_10039.jpg"], ["examples/sa_862.jpg"]]
20
+
21
+ default_example = examples[0]
22
+
23
+ def segment_everything(
24
+ input,
25
+ input_size=1024,
26
+ iou_threshold=0.7,
27
+ conf_threshold=0.25,
28
+ better_quality=False,
29
+ withContours=True,
30
+ use_retina=True,
31
+ text="",
32
+ wider=False,
33
+ mask_random_color=True,
34
+ ):
35
+ input_size = int(input_size)
36
+ w, h = input.size
37
+ scale = input_size / max(w, h)
38
+ new_w = int(w * scale)
39
+ new_h = int(h * scale)
40
+ input = input.resize((new_w, new_h))
41
+
42
+ results = model(input,
43
+ device=device,
44
+ retina_masks=True,
45
+ iou=iou_threshold,
46
+ conf=conf_threshold,
47
+ imgsz=input_size,)
48
+
49
+ if len(text) > 0:
50
+ results = format_results(results[0], 0)
51
+ annotations, _ = text_prompt(results, text, input, device=device, wider=wider)
52
+ annotations = np.array([annotations])
53
+ else:
54
+ annotations = results[0].masks.data
55
+
56
+ fig = fast_process(annotations=annotations,
57
+ image=input,
58
+ device=device,
59
+ scale=(1024 // input_size),
60
+ better_quality=better_quality,
61
+ mask_random_color=mask_random_color,
62
+ bbox=None,
63
+ use_retina=use_retina,
64
+ withContours=withContours,)
65
+ return fig
66
+
67
+
68
+ def segment_with_points(
69
+ input,
70
+ input_size=1024,
71
+ iou_threshold=0.7,
72
+ conf_threshold=0.25,
73
+ better_quality=False,
74
+ withContours=True,
75
+ use_retina=True,
76
+ mask_random_color=True,
77
+ ):
78
+ global global_points
79
+ global global_point_label
80
+
81
+ input_size = int(input_size)
82
+ w, h = input.size
83
+ scale = input_size / max(w, h)
84
+ new_w = int(w * scale)
85
+ new_h = int(h * scale)
86
+ input = input.resize((new_w, new_h))
87
+
88
+ scaled_points = [[int(x * scale) for x in point] for point in global_points]
89
+
90
+ results = model(input,
91
+ device=device,
92
+ retina_masks=True,
93
+ iou=iou_threshold,
94
+ conf=conf_threshold,
95
+ imgsz=input_size,)
96
+
97
+ results = format_results(results[0], 0)
98
+ annotations, _ = point_prompt(results, scaled_points, global_point_label, new_h, new_w)
99
+ annotations = np.array([annotations])
100
+
101
+ fig = fast_process(annotations=annotations,
102
+ image=input,
103
+ device=device,
104
+ scale=(1024 // input_size),
105
+ better_quality=better_quality,
106
+ mask_random_color=mask_random_color,
107
+ bbox=None,
108
+ use_retina=use_retina,
109
+ withContours=withContours,)
110
+
111
+ global_points = []
112
+ global_point_label = []
113
+ return fig, None
114
+
115
+
116
+ def get_points_with_draw(image, label, evt: gr.SelectData):
117
+ global global_points
118
+ global global_point_label
119
+
120
+ x, y = evt.index[0], evt.index[1]
121
+ point_radius, point_color = 15, (255, 255, 0) if label == 'Add Mask' else (255, 0, 255)
122
+ global_points.append([x, y])
123
+ global_point_label.append(1 if label == 'Add Mask' else 0)
124
+
125
+ print(x, y, label == 'Add Mask')
126
+
127
+ draw = ImageDraw.Draw(image)
128
+ draw.ellipse([(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color)
129
+ return image
130
+
131
+
132
+ cond_img_e = gr.Image(label="Input", value=default_example[0], type='pil')
133
+ cond_img_p = gr.Image(label="Input with points", value=default_example[0], type='pil')
134
+ cond_img_t = gr.Image(label="Input with text", value="examples/dogs.jpg", type='pil')
135
+
136
+ segm_img_e = gr.Image(label="Segmented Image", interactive=False, type='pil')
137
+ segm_img_p = gr.Image(label="Segmented Image with points", interactive=False, type='pil')
138
+ segm_img_t = gr.Image(label="Segmented Image with text", interactive=False, type='pil')
139
+
140
+ global_points = []
141
+ global_point_label = []
142
+
143
+ input_size_slider = gr.components.Slider(minimum=512,
144
+ maximum=1024,
145
+ value=1024,
146
+ step=64,
147
+ label='Input_size',
148
+ info='Our model was trained on a size of 1024')
149
+
150
+ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
151
+ with gr.Row():
152
+ with gr.Column(scale=1):
153
+ # Title
154
+ gr.Markdown(title)
155
+
156
+ with gr.Tab("Text mode"):
157
+ # Images
158
+ with gr.Row(variant="panel"):
159
+ with gr.Column(scale=1):
160
+ cond_img_t.render()
161
+
162
+ with gr.Column(scale=1):
163
+ segm_img_t.render()
164
+
165
+ # Submit & Clear
166
+ with gr.Row():
167
+ with gr.Column():
168
+ input_size_slider_t = gr.components.Slider(minimum=512,
169
+ maximum=1024,
170
+ value=1024,
171
+ step=64,
172
+ label='Input_size',
173
+ info='Our model was trained on a size of 1024')
174
+ with gr.Row():
175
+ with gr.Column():
176
+ contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
177
+ text_box = gr.Textbox(label="text prompt", value="a black dog")
178
+
179
+ with gr.Column():
180
+ segment_btn_t = gr.Button("Segment with text", variant='primary')
181
+ clear_btn_t = gr.Button("Clear", variant="secondary")
182
+
183
+ gr.Markdown("Try some of the examples below ⬇️")
184
+ gr.Examples(examples=[["examples/dogs.jpg"], ["examples/fruits.jpg"], ["examples/flowers.jpg"]],
185
+ inputs=[cond_img_t],
186
+ examples_per_page=4)
187
+
188
+ with gr.Column():
189
+ with gr.Accordion("Advanced options", open=False):
190
+ iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
191
+ conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
192
+ with gr.Row():
193
+ mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
194
+ retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
195
+ wider_check = gr.Checkbox(value=False, label='wider', info='wider result')
196
+
197
+ segment_btn_t.click(segment_everything,
198
+ inputs=[
199
+ cond_img_t,
200
+ input_size_slider_t,
201
+ iou_threshold,
202
+ conf_threshold,
203
+ mor_check,
204
+ contour_check,
205
+ retina_check,
206
+ text_box,
207
+ wider_check,
208
+ ],
209
+ outputs=segm_img_t)
210
+
211
+ def clear():
212
+ return None, None
213
+
214
+ def clear_text():
215
+ return None, None, None
216
+
217
+ clear_btn_e.click(clear, outputs=[cond_img_e, segm_img_e])
218
+ clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p])
219
+ clear_btn_t.click(clear_text, outputs=[cond_img_p, segm_img_p, text_box])
220
+
221
+ demo.queue()
222
+ demo.launch()