jameslahm commited on
Commit
719d8e3
1 Parent(s): ee1cac8

Upload 14 files

Browse files
app.py ADDED
@@ -0,0 +1,241 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import gradio as gr
4
+ import numpy as np
5
+ import torch
6
+ from repvit_sam import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry
7
+ from PIL import ImageDraw
8
+ from utils.tools import box_prompt, format_results, point_prompt
9
+ from utils.tools_gradio import fast_process
10
+
11
+ # Most of our demo code is from [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM). Huge thanks for AN-619.
12
+
13
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
+
15
+ # Load the pre-trained model
16
+ sam_checkpoint = "repvit_sam.pt"
17
+ model_type = "repvit"
18
+
19
+ repvit_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
20
+ repvit_sam = repvit_sam.to(device=device)
21
+ repvit_sam.eval()
22
+
23
+ mask_generator = SamAutomaticMaskGenerator(repvit_sam)
24
+ predictor = SamPredictor(repvit_sam)
25
+
26
+ # Description
27
+ title = "<center><strong><font size='8'>RepViT-SAM<font></strong></center>"
28
+
29
+ description_e = """This is a demo of [RepViT-SAM](https://github.com/THU-MIG/RepViT).
30
+
31
+ We will provide box mode soon.
32
+
33
+ Enjoy!
34
+
35
+ """
36
+
37
+ description_p = """ Instructions for point mode
38
+
39
+ 0. Restart by click the Restart button
40
+ 1. Select a point with Add Mask for the foreground (Must)
41
+ 2. Select a point with Remove Area for the background (Optional)
42
+ 3. Click the Start Segmenting.
43
+
44
+ Github [link](https://github.com/THU-MIG/RepViT)
45
+
46
+ """
47
+
48
+ examples = [
49
+ ["app/assets/picture3.jpg"],
50
+ ["app/assets/picture4.jpg"],
51
+ ["app/assets/picture6.jpg"],
52
+ ["app/assets/picture1.jpg"],
53
+ ]
54
+
55
+ default_example = examples[0]
56
+
57
+ css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
58
+
59
+ def segment_with_points(
60
+ image,
61
+ original_image,
62
+ input_size=1024,
63
+ better_quality=False,
64
+ withContours=True,
65
+ use_retina=True,
66
+ mask_random_color=True,
67
+ ):
68
+ global global_points
69
+ global global_point_label
70
+
71
+ input_size = int(input_size)
72
+ w, h = image.size
73
+ scale = input_size / max(w, h)
74
+ new_w = int(w * scale)
75
+ new_h = int(h * scale)
76
+ image = image.resize((new_w, new_h))
77
+
78
+ scaled_points = np.array(
79
+ [[int(x * scale) for x in point] for point in global_points]
80
+ )
81
+ scaled_point_label = np.array(global_point_label)
82
+
83
+ if scaled_points.size == 0 and scaled_point_label.size == 0:
84
+ print("No points selected")
85
+ return image, image
86
+
87
+ nd_image = np.array(original_image.resize((new_w, new_h)))
88
+ predictor.set_image(nd_image)
89
+ masks, scores, logits = predictor.predict(
90
+ point_coords=scaled_points,
91
+ point_labels=scaled_point_label,
92
+ multimask_output=False,
93
+ )
94
+
95
+ results = format_results(masks, scores, logits, 0)
96
+
97
+ annotations, _ = point_prompt(
98
+ results, scaled_points, scaled_point_label, new_h, new_w
99
+ )
100
+ annotations = np.array([annotations])
101
+
102
+ fig = fast_process(
103
+ annotations=annotations,
104
+ image=image,
105
+ device=device,
106
+ scale=(1024 // input_size),
107
+ better_quality=better_quality,
108
+ mask_random_color=mask_random_color,
109
+ bbox=None,
110
+ use_retina=use_retina,
111
+ withContours=withContours,
112
+ )
113
+
114
+ global_points = []
115
+ global_point_label = []
116
+ # return fig, None
117
+ return fig, original_image.resize((new_w, new_h))
118
+
119
+
120
+ def get_points_with_draw(image, label, evt: gr.SelectData):
121
+ global global_points
122
+ global global_point_label
123
+
124
+ x, y = evt.index[0], evt.index[1]
125
+ point_radius, point_color = 15 * ((max(image.width, image.height)) / 1024), (255, 255, 0) if label == "Add Mask" else (
126
+ 255,
127
+ 0,
128
+ 255,
129
+ )
130
+ global_points.append([x, y])
131
+ global_point_label.append(1 if label == "Add Mask" else 0)
132
+
133
+ # 创建一个可以在图像上绘图的对象
134
+ draw = ImageDraw.Draw(image)
135
+ draw.ellipse(
136
+ [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
137
+ fill=point_color,
138
+ )
139
+ return image
140
+
141
+
142
+ cond_img_e = gr.Image(label="Input", value=default_example[0], type="pil")
143
+ cond_img_p = gr.Image(label="Input with points", value=default_example[0], type="pil")
144
+
145
+ segm_img_e = gr.Image(label="Segmented Image", interactive=False, type="pil")
146
+ segm_img_p = gr.Image(
147
+ label="Segmented Image with points", interactive=True, type="pil"
148
+ )
149
+
150
+ global_points = []
151
+ global_point_label = []
152
+
153
+ input_size_slider = gr.components.Slider(
154
+ minimum=512,
155
+ maximum=1024,
156
+ value=1024,
157
+ step=64,
158
+ label="Input_size",
159
+ info="Our model was trained on a size of 1024",
160
+ )
161
+
162
+ with gr.Blocks(css=css, title="RepViT-SAM") as demo:
163
+ from PIL import Image
164
+ original_image = gr.State(value=Image.open(default_example[0]).convert('RGB'))
165
+
166
+ with gr.Row():
167
+ with gr.Column(scale=1):
168
+ # Title
169
+ gr.Markdown(title)
170
+
171
+ with gr.Tab("Point mode"):
172
+ # Images
173
+ with gr.Row(variant="panel"):
174
+ with gr.Column(scale=1):
175
+ cond_img_p.render()
176
+
177
+ with gr.Column(scale=1):
178
+ segm_img_p.render()
179
+
180
+ # Submit & Clear
181
+ with gr.Row():
182
+ with gr.Column():
183
+ with gr.Row():
184
+ add_or_remove = gr.Radio(
185
+ ["Add Mask", "Remove Area"],
186
+ value="Add Mask",
187
+ )
188
+
189
+ with gr.Column():
190
+ segment_btn_p = gr.Button(
191
+ "Start segmenting!", variant="primary"
192
+ )
193
+ clear_btn_p = gr.Button("Restart", variant="secondary")
194
+
195
+ gr.Markdown("Try some of the examples below ⬇️")
196
+
197
+ gr.Examples(
198
+ examples=examples,
199
+ inputs=[cond_img_p],
200
+ fn=lambda x: x,
201
+ outputs=[original_image],
202
+ # fn=segment_with_points,
203
+ # cache_examples=True,
204
+ examples_per_page=4,
205
+ run_on_click=True
206
+ )
207
+
208
+ with gr.Column():
209
+ # Description
210
+ gr.Markdown(description_p)
211
+
212
+ cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)
213
+ cond_img_p.upload(lambda x: x, inputs=[cond_img_p], outputs=[original_image])
214
+
215
+ # segment_btn_e.click(
216
+ # segment_everything,
217
+ # inputs=[
218
+ # cond_img_e,
219
+ # input_size_slider,
220
+ # mor_check,
221
+ # contour_check,
222
+ # retina_check,
223
+ # ],
224
+ # outputs=segm_img_e,
225
+ # )
226
+
227
+ segment_btn_p.click(
228
+ segment_with_points, inputs=[cond_img_p, original_image], outputs=[segm_img_p, cond_img_p]
229
+ )
230
+
231
+ def clear():
232
+ return None, None
233
+
234
+ def clear_text():
235
+ return None, None, None
236
+
237
+ # clear_btn_e.click(clear, outputs=[cond_img_e, segm_img_e])
238
+ clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p])
239
+
240
+ demo.queue()
241
+ demo.launch()
assets/picture1.jpg ADDED
assets/picture2.jpg ADDED
assets/picture3.jpg ADDED
assets/picture4.jpg ADDED
assets/picture5.jpg ADDED
assets/picture6.jpg ADDED
repvit_sam.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1fd2867cf669e6d89b1c91d0eb883d83ef4bcc590143dd2ca35e2705b21a4f2f
3
+ size 110032539
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ timm
4
+ opencv-python
5
+ git+https://git@github.com/THU-RepViT/RepViT.git#subdirectory=sam
utils/__init__.py ADDED
File without changes
utils/tools.py ADDED
@@ -0,0 +1,406 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+
4
+ import cv2
5
+ import matplotlib.pyplot as plt
6
+ import numpy as np
7
+ import torch
8
+ from PIL import Image
9
+
10
+
11
+ def convert_box_xywh_to_xyxy(box):
12
+ x1 = box[0]
13
+ y1 = box[1]
14
+ x2 = box[0] + box[2]
15
+ y2 = box[1] + box[3]
16
+ return [x1, y1, x2, y2]
17
+
18
+
19
+ def segment_image(image, bbox):
20
+ image_array = np.array(image)
21
+ segmented_image_array = np.zeros_like(image_array)
22
+ x1, y1, x2, y2 = bbox
23
+ segmented_image_array[y1:y2, x1:x2] = image_array[y1:y2, x1:x2]
24
+ segmented_image = Image.fromarray(segmented_image_array)
25
+ black_image = Image.new("RGB", image.size, (255, 255, 255))
26
+ # transparency_mask = np.zeros_like((), dtype=np.uint8)
27
+ transparency_mask = np.zeros(
28
+ (image_array.shape[0], image_array.shape[1]), dtype=np.uint8
29
+ )
30
+ transparency_mask[y1:y2, x1:x2] = 255
31
+ transparency_mask_image = Image.fromarray(transparency_mask, mode="L")
32
+ black_image.paste(segmented_image, mask=transparency_mask_image)
33
+ return black_image
34
+
35
+
36
+ def format_results(masks, scores, logits, filter=0):
37
+ annotations = []
38
+ n = len(scores)
39
+ for i in range(n):
40
+ annotation = {}
41
+
42
+ mask = masks[i]
43
+ tmp = np.where(mask != 0)
44
+ if np.sum(mask) < filter:
45
+ continue
46
+ annotation["id"] = i
47
+ annotation["segmentation"] = mask
48
+ annotation["bbox"] = [
49
+ np.min(tmp[0]),
50
+ np.min(tmp[1]),
51
+ np.max(tmp[1]),
52
+ np.max(tmp[0]),
53
+ ]
54
+ annotation["score"] = scores[i]
55
+ annotation["area"] = annotation["segmentation"].sum()
56
+ annotations.append(annotation)
57
+ return annotations
58
+
59
+
60
+ def filter_masks(annotations): # filter the overlap mask
61
+ annotations.sort(key=lambda x: x["area"], reverse=True)
62
+ to_remove = set()
63
+ for i in range(0, len(annotations)):
64
+ a = annotations[i]
65
+ for j in range(i + 1, len(annotations)):
66
+ b = annotations[j]
67
+ if i != j and j not in to_remove:
68
+ # check if
69
+ if b["area"] < a["area"]:
70
+ if (a["segmentation"] & b["segmentation"]).sum() / b[
71
+ "segmentation"
72
+ ].sum() > 0.8:
73
+ to_remove.add(j)
74
+
75
+ return [a for i, a in enumerate(annotations) if i not in to_remove], to_remove
76
+
77
+
78
+ def get_bbox_from_mask(mask):
79
+ mask = mask.astype(np.uint8)
80
+ contours, hierarchy = cv2.findContours(
81
+ mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
82
+ )
83
+ x1, y1, w, h = cv2.boundingRect(contours[0])
84
+ x2, y2 = x1 + w, y1 + h
85
+ if len(contours) > 1:
86
+ for b in contours:
87
+ x_t, y_t, w_t, h_t = cv2.boundingRect(b)
88
+ # 将多个bbox合并成一个
89
+ x1 = min(x1, x_t)
90
+ y1 = min(y1, y_t)
91
+ x2 = max(x2, x_t + w_t)
92
+ y2 = max(y2, y_t + h_t)
93
+ h = y2 - y1
94
+ w = x2 - x1
95
+ return [x1, y1, x2, y2]
96
+
97
+
98
+ def fast_process(
99
+ annotations, args, mask_random_color, bbox=None, points=None, edges=False
100
+ ):
101
+ if isinstance(annotations[0], dict):
102
+ annotations = [annotation["segmentation"] for annotation in annotations]
103
+ result_name = os.path.basename(args.img_path)
104
+ image = cv2.imread(args.img_path)
105
+ image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
106
+ original_h = image.shape[0]
107
+ original_w = image.shape[1]
108
+ if sys.platform == "darwin":
109
+ plt.switch_backend("TkAgg")
110
+ plt.figure(figsize=(original_w / 100, original_h / 100))
111
+ # Add subplot with no margin.
112
+ plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
113
+ plt.margins(0, 0)
114
+ plt.gca().xaxis.set_major_locator(plt.NullLocator())
115
+ plt.gca().yaxis.set_major_locator(plt.NullLocator())
116
+ plt.imshow(image)
117
+ if args.better_quality == True:
118
+ if isinstance(annotations[0], torch.Tensor):
119
+ annotations = np.array(annotations.cpu())
120
+ for i, mask in enumerate(annotations):
121
+ mask = cv2.morphologyEx(
122
+ mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
123
+ )
124
+ annotations[i] = cv2.morphologyEx(
125
+ mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
126
+ )
127
+ if args.device == "cpu":
128
+ annotations = np.array(annotations)
129
+ fast_show_mask(
130
+ annotations,
131
+ plt.gca(),
132
+ random_color=mask_random_color,
133
+ bbox=bbox,
134
+ points=points,
135
+ point_label=args.point_label,
136
+ retinamask=args.retina,
137
+ target_height=original_h,
138
+ target_width=original_w,
139
+ )
140
+ else:
141
+ if isinstance(annotations[0], np.ndarray):
142
+ annotations = torch.from_numpy(annotations)
143
+ fast_show_mask_gpu(
144
+ annotations,
145
+ plt.gca(),
146
+ random_color=args.randomcolor,
147
+ bbox=bbox,
148
+ points=points,
149
+ point_label=args.point_label,
150
+ retinamask=args.retina,
151
+ target_height=original_h,
152
+ target_width=original_w,
153
+ )
154
+ if isinstance(annotations, torch.Tensor):
155
+ annotations = annotations.cpu().numpy()
156
+ if args.withContours == True:
157
+ contour_all = []
158
+ temp = np.zeros((original_h, original_w, 1))
159
+ for i, mask in enumerate(annotations):
160
+ if type(mask) == dict:
161
+ mask = mask["segmentation"]
162
+ annotation = mask.astype(np.uint8)
163
+ if args.retina == False:
164
+ annotation = cv2.resize(
165
+ annotation,
166
+ (original_w, original_h),
167
+ interpolation=cv2.INTER_NEAREST,
168
+ )
169
+ contours, hierarchy = cv2.findContours(
170
+ annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
171
+ )
172
+ for contour in contours:
173
+ contour_all.append(contour)
174
+ cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
175
+ color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
176
+ contour_mask = temp / 255 * color.reshape(1, 1, -1)
177
+ plt.imshow(contour_mask)
178
+
179
+ save_path = args.output
180
+ if not os.path.exists(save_path):
181
+ os.makedirs(save_path)
182
+ plt.axis("off")
183
+ fig = plt.gcf()
184
+ plt.draw()
185
+
186
+ try:
187
+ buf = fig.canvas.tostring_rgb()
188
+ except AttributeError:
189
+ fig.canvas.draw()
190
+ buf = fig.canvas.tostring_rgb()
191
+
192
+ cols, rows = fig.canvas.get_width_height()
193
+ img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
194
+ cv2.imwrite(
195
+ os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
196
+ )
197
+
198
+
199
+ # CPU post process
200
+ def fast_show_mask(
201
+ annotation,
202
+ ax,
203
+ random_color=False,
204
+ bbox=None,
205
+ points=None,
206
+ point_label=None,
207
+ retinamask=True,
208
+ target_height=960,
209
+ target_width=960,
210
+ ):
211
+ msak_sum = annotation.shape[0]
212
+ height = annotation.shape[1]
213
+ weight = annotation.shape[2]
214
+ # 将annotation 按照面积 排序
215
+ areas = np.sum(annotation, axis=(1, 2))
216
+ sorted_indices = np.argsort(areas)
217
+ annotation = annotation[sorted_indices]
218
+
219
+ index = (annotation != 0).argmax(axis=0)
220
+ if random_color == True:
221
+ color = np.random.random((msak_sum, 1, 1, 3))
222
+ else:
223
+ color = np.ones((msak_sum, 1, 1, 3)) * np.array(
224
+ [30 / 255, 144 / 255, 255 / 255]
225
+ )
226
+ transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
227
+ visual = np.concatenate([color, transparency], axis=-1)
228
+ mask_image = np.expand_dims(annotation, -1) * visual
229
+
230
+ show = np.zeros((height, weight, 4))
231
+ h_indices, w_indices = np.meshgrid(
232
+ np.arange(height), np.arange(weight), indexing="ij"
233
+ )
234
+ indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
235
+ # 使用向量化索引更新show的值
236
+ show[h_indices, w_indices, :] = mask_image[indices]
237
+ if bbox is not None:
238
+ x1, y1, x2, y2 = bbox
239
+ ax.add_patch(
240
+ plt.Rectangle(
241
+ (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
242
+ )
243
+ )
244
+ # draw point
245
+ if points is not None:
246
+ plt.scatter(
247
+ [point[0] for i, point in enumerate(points) if point_label[i] == 1],
248
+ [point[1] for i, point in enumerate(points) if point_label[i] == 1],
249
+ s=20,
250
+ c="y",
251
+ )
252
+ plt.scatter(
253
+ [point[0] for i, point in enumerate(points) if point_label[i] == 0],
254
+ [point[1] for i, point in enumerate(points) if point_label[i] == 0],
255
+ s=20,
256
+ c="m",
257
+ )
258
+
259
+ if retinamask == False:
260
+ show = cv2.resize(
261
+ show, (target_width, target_height), interpolation=cv2.INTER_NEAREST
262
+ )
263
+ ax.imshow(show)
264
+
265
+
266
+ def fast_show_mask_gpu(
267
+ annotation,
268
+ ax,
269
+ random_color=False,
270
+ bbox=None,
271
+ points=None,
272
+ point_label=None,
273
+ retinamask=True,
274
+ target_height=960,
275
+ target_width=960,
276
+ ):
277
+ msak_sum = annotation.shape[0]
278
+ height = annotation.shape[1]
279
+ weight = annotation.shape[2]
280
+ areas = torch.sum(annotation, dim=(1, 2))
281
+ sorted_indices = torch.argsort(areas, descending=False)
282
+ annotation = annotation[sorted_indices]
283
+ # 找每个位置第一个非零值下标
284
+ index = (annotation != 0).to(torch.long).argmax(dim=0)
285
+ if random_color == True:
286
+ color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
287
+ else:
288
+ color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor(
289
+ [30 / 255, 144 / 255, 255 / 255]
290
+ ).to(annotation.device)
291
+ transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
292
+ visual = torch.cat([color, transparency], dim=-1)
293
+ mask_image = torch.unsqueeze(annotation, -1) * visual
294
+ # 按index取数,index指每个位置选哪个batch的数,把mask_image转成一个batch的形式
295
+ show = torch.zeros((height, weight, 4)).to(annotation.device)
296
+ h_indices, w_indices = torch.meshgrid(
297
+ torch.arange(height), torch.arange(weight), indexing="ij"
298
+ )
299
+ indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
300
+ # 使用向量化索引更新show的值
301
+ show[h_indices, w_indices, :] = mask_image[indices]
302
+ show_cpu = show.cpu().numpy()
303
+ if bbox is not None:
304
+ x1, y1, x2, y2 = bbox
305
+ ax.add_patch(
306
+ plt.Rectangle(
307
+ (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
308
+ )
309
+ )
310
+ # draw point
311
+ if points is not None:
312
+ plt.scatter(
313
+ [point[0] for i, point in enumerate(points) if point_label[i] == 1],
314
+ [point[1] for i, point in enumerate(points) if point_label[i] == 1],
315
+ s=20,
316
+ c="y",
317
+ )
318
+ plt.scatter(
319
+ [point[0] for i, point in enumerate(points) if point_label[i] == 0],
320
+ [point[1] for i, point in enumerate(points) if point_label[i] == 0],
321
+ s=20,
322
+ c="m",
323
+ )
324
+ if retinamask == False:
325
+ show_cpu = cv2.resize(
326
+ show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
327
+ )
328
+ ax.imshow(show_cpu)
329
+
330
+
331
+ def crop_image(annotations, image_like):
332
+ if isinstance(image_like, str):
333
+ image = Image.open(image_like)
334
+ else:
335
+ image = image_like
336
+ ori_w, ori_h = image.size
337
+ mask_h, mask_w = annotations[0]["segmentation"].shape
338
+ if ori_w != mask_w or ori_h != mask_h:
339
+ image = image.resize((mask_w, mask_h))
340
+ cropped_boxes = []
341
+ cropped_images = []
342
+ not_crop = []
343
+ filter_id = []
344
+ # annotations, _ = filter_masks(annotations)
345
+ # filter_id = list(_)
346
+ for _, mask in enumerate(annotations):
347
+ if np.sum(mask["segmentation"]) <= 100:
348
+ filter_id.append(_)
349
+ continue
350
+ bbox = get_bbox_from_mask(mask["segmentation"]) # mask 的 bbox
351
+ cropped_boxes.append(segment_image(image, bbox)) # 保存裁剪的图片
352
+ # cropped_boxes.append(segment_image(image,mask["segmentation"]))
353
+ cropped_images.append(bbox) # 保存裁剪的图片的bbox
354
+
355
+ return cropped_boxes, cropped_images, not_crop, filter_id, annotations
356
+
357
+
358
+ def box_prompt(masks, bbox, target_height, target_width):
359
+ h = masks.shape[1]
360
+ w = masks.shape[2]
361
+ if h != target_height or w != target_width:
362
+ bbox = [
363
+ int(bbox[0] * w / target_width),
364
+ int(bbox[1] * h / target_height),
365
+ int(bbox[2] * w / target_width),
366
+ int(bbox[3] * h / target_height),
367
+ ]
368
+ bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
369
+ bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
370
+ bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
371
+ bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h
372
+
373
+ # IoUs = torch.zeros(len(masks), dtype=torch.float32)
374
+ bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
375
+
376
+ masks_area = torch.sum(masks[:, bbox[1] : bbox[3], bbox[0] : bbox[2]], dim=(1, 2))
377
+ orig_masks_area = torch.sum(masks, dim=(1, 2))
378
+
379
+ union = bbox_area + orig_masks_area - masks_area
380
+ IoUs = masks_area / union
381
+ max_iou_index = torch.argmax(IoUs)
382
+
383
+ return masks[max_iou_index].cpu().numpy(), max_iou_index
384
+
385
+
386
+ def point_prompt(masks, points, point_label, target_height, target_width): # numpy 处理
387
+ h = masks[0]["segmentation"].shape[0]
388
+ w = masks[0]["segmentation"].shape[1]
389
+ if h != target_height or w != target_width:
390
+ points = [
391
+ [int(point[0] * w / target_width), int(point[1] * h / target_height)]
392
+ for point in points
393
+ ]
394
+ onemask = np.zeros((h, w))
395
+ for i, annotation in enumerate(masks):
396
+ if type(annotation) == dict:
397
+ mask = annotation["segmentation"]
398
+ else:
399
+ mask = annotation
400
+ for i, point in enumerate(points):
401
+ if mask[point[1], point[0]] == 1 and point_label[i] == 1:
402
+ onemask += mask
403
+ if mask[point[1], point[0]] == 1 and point_label[i] == 0:
404
+ onemask -= mask
405
+ onemask = onemask >= 1
406
+ return onemask, 0
utils/tools_gradio.py ADDED
@@ -0,0 +1,192 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import matplotlib.pyplot as plt
3
+ import numpy as np
4
+ import torch
5
+ from PIL import Image
6
+
7
+
8
+ def fast_process(
9
+ annotations,
10
+ image,
11
+ device,
12
+ scale,
13
+ better_quality=False,
14
+ mask_random_color=True,
15
+ bbox=None,
16
+ use_retina=True,
17
+ withContours=True,
18
+ ):
19
+ if isinstance(annotations[0], dict):
20
+ annotations = [annotation["segmentation"] for annotation in annotations]
21
+
22
+ original_h = image.height
23
+ original_w = image.width
24
+ if better_quality:
25
+ if isinstance(annotations[0], torch.Tensor):
26
+ annotations = np.array(annotations.cpu())
27
+ for i, mask in enumerate(annotations):
28
+ mask = cv2.morphologyEx(
29
+ mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
30
+ )
31
+ annotations[i] = cv2.morphologyEx(
32
+ mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
33
+ )
34
+ if device == "cpu":
35
+ annotations = np.array(annotations)
36
+ inner_mask = fast_show_mask(
37
+ annotations,
38
+ plt.gca(),
39
+ random_color=mask_random_color,
40
+ bbox=bbox,
41
+ retinamask=use_retina,
42
+ target_height=original_h,
43
+ target_width=original_w,
44
+ )
45
+ else:
46
+ if isinstance(annotations[0], np.ndarray):
47
+ annotations = np.array(annotations)
48
+ annotations = torch.from_numpy(annotations)
49
+ inner_mask = fast_show_mask_gpu(
50
+ annotations,
51
+ plt.gca(),
52
+ random_color=mask_random_color,
53
+ bbox=bbox,
54
+ retinamask=use_retina,
55
+ target_height=original_h,
56
+ target_width=original_w,
57
+ )
58
+ if isinstance(annotations, torch.Tensor):
59
+ annotations = annotations.cpu().numpy()
60
+
61
+ if withContours:
62
+ contour_all = []
63
+ temp = np.zeros((original_h, original_w, 1))
64
+ for i, mask in enumerate(annotations):
65
+ if type(mask) == dict:
66
+ mask = mask["segmentation"]
67
+ annotation = mask.astype(np.uint8)
68
+ if use_retina == False:
69
+ annotation = cv2.resize(
70
+ annotation,
71
+ (original_w, original_h),
72
+ interpolation=cv2.INTER_NEAREST,
73
+ )
74
+ contours, _ = cv2.findContours(
75
+ annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
76
+ )
77
+ for contour in contours:
78
+ contour_all.append(contour)
79
+ cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2 // scale)
80
+ color = np.array([0 / 255, 0 / 255, 255 / 255, 0.9])
81
+ contour_mask = temp / 255 * color.reshape(1, 1, -1)
82
+
83
+ image = image.convert("RGBA")
84
+ overlay_inner = Image.fromarray((inner_mask * 255).astype(np.uint8), "RGBA")
85
+ image.paste(overlay_inner, (0, 0), overlay_inner)
86
+
87
+ if withContours:
88
+ overlay_contour = Image.fromarray((contour_mask * 255).astype(np.uint8), "RGBA")
89
+ image.paste(overlay_contour, (0, 0), overlay_contour)
90
+
91
+ return image
92
+
93
+
94
+ # CPU post process
95
+ def fast_show_mask(
96
+ annotation,
97
+ ax,
98
+ random_color=False,
99
+ bbox=None,
100
+ retinamask=True,
101
+ target_height=960,
102
+ target_width=960,
103
+ ):
104
+ mask_sum = annotation.shape[0]
105
+ height = annotation.shape[1]
106
+ weight = annotation.shape[2]
107
+ # 将annotation 按照面积 排序
108
+ areas = np.sum(annotation, axis=(1, 2))
109
+ sorted_indices = np.argsort(areas)[::1]
110
+ annotation = annotation[sorted_indices]
111
+
112
+ index = (annotation != 0).argmax(axis=0)
113
+ if random_color == True:
114
+ color = np.random.random((mask_sum, 1, 1, 3))
115
+ else:
116
+ color = np.ones((mask_sum, 1, 1, 3)) * np.array(
117
+ [30 / 255, 144 / 255, 255 / 255]
118
+ )
119
+ transparency = np.ones((mask_sum, 1, 1, 1)) * 0.6
120
+ visual = np.concatenate([color, transparency], axis=-1)
121
+ mask_image = np.expand_dims(annotation, -1) * visual
122
+
123
+ mask = np.zeros((height, weight, 4))
124
+
125
+ h_indices, w_indices = np.meshgrid(
126
+ np.arange(height), np.arange(weight), indexing="ij"
127
+ )
128
+ indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
129
+
130
+ mask[h_indices, w_indices, :] = mask_image[indices]
131
+ if bbox is not None:
132
+ x1, y1, x2, y2 = bbox
133
+ ax.add_patch(
134
+ plt.Rectangle(
135
+ (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
136
+ )
137
+ )
138
+
139
+ if retinamask == False:
140
+ mask = cv2.resize(
141
+ mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST
142
+ )
143
+
144
+ return mask
145
+
146
+
147
+ def fast_show_mask_gpu(
148
+ annotation,
149
+ ax,
150
+ random_color=False,
151
+ bbox=None,
152
+ retinamask=True,
153
+ target_height=960,
154
+ target_width=960,
155
+ ):
156
+ device = annotation.device
157
+ mask_sum = annotation.shape[0]
158
+ height = annotation.shape[1]
159
+ weight = annotation.shape[2]
160
+ areas = torch.sum(annotation, dim=(1, 2))
161
+ sorted_indices = torch.argsort(areas, descending=False)
162
+ annotation = annotation[sorted_indices]
163
+ # 找每个位置第一个非零值下标
164
+ index = (annotation != 0).to(torch.long).argmax(dim=0)
165
+ if random_color == True:
166
+ color = torch.rand((mask_sum, 1, 1, 3)).to(device)
167
+ else:
168
+ color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
169
+ [30 / 255, 144 / 255, 255 / 255]
170
+ ).to(device)
171
+ transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
172
+ visual = torch.cat([color, transparency], dim=-1)
173
+ mask_image = torch.unsqueeze(annotation, -1) * visual
174
+ # 按index取数,index指每个位置选哪个batch的数,把mask_image转成一个batch的形式
175
+ mask = torch.zeros((height, weight, 4)).to(device)
176
+ h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
177
+ indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
178
+ # 使用向量化索引更新show的值
179
+ mask[h_indices, w_indices, :] = mask_image[indices]
180
+ mask_cpu = mask.cpu().numpy()
181
+ if bbox is not None:
182
+ x1, y1, x2, y2 = bbox
183
+ ax.add_patch(
184
+ plt.Rectangle(
185
+ (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
186
+ )
187
+ )
188
+ if retinamask == False:
189
+ mask_cpu = cv2.resize(
190
+ mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
191
+ )
192
+ return mask_cpu