dragonSwing/annotate-anything
2
1import functools
2import json
3import os
4import sys
5import tempfile
6
7import cv2
8import gradio as gr
9import numpy as np
10import supervision as sv
11import torch
12from PIL import Image
13from segment_anything import build_sam
14from segment_anything import SamAutomaticMaskGenerator
15from segment_anything import SamPredictor
16from supervision.detection.utils import mask_to_polygons
17from supervision.detection.utils import xywh_to_xyxy
18
19if os.environ.get("IS_MY_DEBUG") is None:
20 os.system("pip install -e GroundingDINO")
21
22sys.path.append("tag2text")
23sys.path.append("GroundingDINO")
24
25from groundingdino.util.inference import Model as DinoModel
26from tag2text.models import tag2text
27from config import *
28from utils import download_file_hf, detect, segment, generate_tags
29
30if not os.path.exists(abs_weight_dir):
31 os.makedirs(abs_weight_dir, exist_ok=True)
32
33sam_checkpoint = os.path.join(abs_weight_dir, sam_dict[default_sam]["checkpoint_file"])
34if not os.path.exists(sam_checkpoint):
35 os.system(f"wget {sam_dict[default_sam]['checkpoint_url']} -O {sam_checkpoint}")
36
37tag2text_checkpoint = os.path.join(
38 abs_weight_dir, tag2text_dict[default_tag2text]["checkpoint_file"]
39)
40if not os.path.exists(tag2text_checkpoint):
41 os.system(
42 f"wget {tag2text_dict[default_tag2text]['checkpoint_url']} -O {tag2text_checkpoint}"
43 )
44
45dino_checkpoint = os.path.join(
46 abs_weight_dir, dino_dict[default_dino]["checkpoint_file"]
47)
48dino_config_file = os.path.join(abs_weight_dir, dino_dict[default_dino]["config_file"])
49if not os.path.exists(dino_checkpoint):
50 dino_repo_id = dino_dict[default_dino]["repo_id"]
51 download_file_hf(
52 repo_id=dino_repo_id,
53 filename=dino_dict[default_dino]["config_file"],
54 cache_dir=weight_dir,
55 )
56 download_file_hf(
57 repo_id=dino_repo_id,
58 filename=dino_dict[default_dino]["checkpoint_file"],
59 cache_dir=weight_dir,
60 )
61
62# load model
63device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
64tag2text_model = tag2text.tag2text_caption(
65 pretrained=tag2text_checkpoint,
66 image_size=384,
67 vit="swin_b",
68 delete_tag_index=delete_tag_index,
69)
70# threshold for tagging
71# we reduce the threshold to obtain more tags
72tag2text_model.threshold = 0.64
73tag2text_model.to(device)
74tag2text_model.eval()
75
76
77sam = build_sam(checkpoint=sam_checkpoint)
78sam.to(device=device)
79sam_predictor = SamPredictor(sam)
80sam_automask_generator = SamAutomaticMaskGenerator(sam)
81
82grounding_dino_model = DinoModel(
83 model_config_path=dino_config_file,
84 model_checkpoint_path=dino_checkpoint,
85 device=device,
86)
87
88
89def process(
90 image_path,
91 task,
92 prompt,
93 box_threshold,
94 text_threshold,
95 iou_threshold,
96 kernel_size,
97 expand_mask,
98):
99 global tag2text_model, sam_predictor, sam_automask_generator, grounding_dino_model, device
100 output_gallery = []
101 detections = None
102 metadata = {"image": {}, "annotations": []}
103
104 try:
105 # Load image
106 image = Image.open(image_path)
107 image_pil = image.convert("RGB")
108 image = np.array(image_pil)
109 orig_image = image.copy()
110
111 # Extract image metadata
112 filename = os.path.basename(image_path)
113 h, w = image.shape[:2]
114 metadata["image"]["file_name"] = filename
115 metadata["image"]["width"] = w
116 metadata["image"]["height"] = h
117
118 # Generate tags
119 if task in ["auto", "detection"] and prompt == "":
120 tags, caption = generate_tags(tag2text_model, image_pil, "None", device)
121 prompt = " . ".join(tags)
122 print(f"Caption: {caption}")
123 print(f"Tags: {tags}")
124
125 # ToDo: Extract metadata
126 metadata["image"]["caption"] = caption
127 metadata["image"]["tags"] = tags
128
129 if prompt:
130 metadata["prompt"] = prompt
131 print(f"Prompt: {prompt}")
132
133 # Detect boxes
134 if prompt != "":
135 detections, phrases, classes = detect(
136 grounding_dino_model,
137 image,
138 caption=prompt,
139 box_threshold=box_threshold,
140 text_threshold=text_threshold,
141 iou_threshold=iou_threshold,
142 post_process=True,
143 )
144 print(phrases)
145
146 # Draw boxes
147 box_annotator = sv.BoxAnnotator()
148 labels = [
149 f"{phrases[i]} {detections.confidence[i]:0.2f}"
150 for i in range(len(phrases))
151 ]
152 image = box_annotator.annotate(
153 scene=image, detections=detections, labels=labels
154 )
155 output_gallery.append(image)
156
157 # Segmentation
158 if task in ["auto", "segment"]:
159 kernel = cv2.getStructuringElement(
160 cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
161 )
162 if detections:
163 masks, scores = segment(
164 sam_predictor, image=orig_image, boxes=detections.xyxy
165 )
166 if expand_mask:
167 masks = [
168 cv2.dilate(mask.astype(np.uint8), kernel) for mask in masks
169 ]
170 else:
171 masks = [
172 cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
173 for mask in masks
174 ]
175 detections.mask = masks
176 binary_mask = functools.reduce(
177 lambda x, y: x + y, detections.mask
178 ).astype(bool)
179 else:
180 masks = sam_automask_generator.generate(orig_image)
181 sorted_generated_masks = sorted(
182 masks, key=lambda x: x["area"], reverse=True
183 )
184
185 xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
186 scores = np.array(
187 [mask["predicted_iou"] for mask in sorted_generated_masks]
188 )
189 if expand_mask:
190 mask = np.array(
191 [
192 cv2.dilate(mask["segmentation"].astype(np.uint8), kernel)
193 for mask in sorted_generated_masks
194 ]
195 )
196 else:
197 mask = np.array(
198 [mask["segmentation"] for mask in sorted_generated_masks]
199 )
200 detections = sv.Detections(
201 xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask
202 )
203 binary_mask = None
204
205 mask_annotator = sv.MaskAnnotator()
206 mask_image = np.zeros_like(image, dtype=np.uint8)
207 mask_image = mask_annotator.annotate(
208 mask_image, detections=detections, opacity=1
209 )
210 annotated_image = mask_annotator.annotate(image, detections=detections)
211
212 output_gallery.append(mask_image)
213 if binary_mask is not None:
214 binary_mask_image = binary_mask * 255
215 cutout_image = np.expand_dims(binary_mask, axis=-1) * orig_image
216 output_gallery.append(binary_mask_image)
217 output_gallery.append(cutout_image)
218 output_gallery.append(annotated_image)
219
220 # ToDo: Extract metadata
221 if detections:
222 i = 0
223 for (xyxy, mask, confidence, _, _), area, box_area in zip(
224 detections, detections.area, detections.box_area
225 ):
226 annotation = {
227 "id": i + 1,
228 "bbox": [int(x) for x in xyxy],
229 "box_area": float(box_area),
230 }
231 if confidence:
232 annotation["confidence"] = float(confidence)
233 annotation["label"] = phrases[i]
234 if mask is not None:
235 # annotation["segmentation"] = mask_to_polygons(mask)
236 annotation["area"] = int(area)
237 annotation["predicted_iou"] = float(scores[i])
238 metadata["annotations"].append(annotation)
239 i += 1
240
241 meta_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
242 meta_file_path = meta_file.name
243 with open(meta_file_path, "w", encoding="utf-8") as fp:
244 json.dump(metadata, fp)
245
246 return output_gallery, meta_file_path
247 except Exception as error:
248 raise gr.Error(f"global exception: {error}")
249
250
251title = "Annotate Anything"
252
253with gr.Blocks(css="style.css", title=title) as demo:
254 with gr.Row(elem_classes=["container"]):
255 with gr.Column(scale=1):
256 input_image = gr.Image(type="filepath", label="Input")
257 task = gr.Dropdown(
258 ["detect", "segment", "auto"], value="auto", label="task_type"
259 )
260 text_prompt = gr.Textbox(
261 label="Detection Prompt",
262 info="To detect multiple objects, seperating each name with '.', like this: cat . dog . chair ",
263 )
264 with gr.Accordion("Advanced parameters", open=False):
265 box_threshold = gr.Slider(
266 minimum=0,
267 maximum=1,
268 value=0.3,
269 step=0.05,
270 label="Box threshold",
271 )
272 text_threshold = gr.Slider(
273 minimum=0,
274 maximum=1,
275 value=0.25,
276 step=0.05,
277 label="Text threshold",
278 )
279 iou_threshold = gr.Slider(
280 minimum=0,
281 maximum=1,
282 value=0.5,
283 step=0.05,
284 label="IOU threshold",
285 info="Intersection over Union threshold",
286 )
287 kernel_size = gr.Slider(
288 minimum=1,
289 maximum=5,
290 value=2,
291 step=1,
292 label="Kernel size",
293 info="Use to smooth segment masks",
294 )
295 expand_mask = gr.Checkbox(
296 label="Expand mask",
297 )
298 run_button = gr.Button(label="Run")
299
300 with gr.Column(scale=2):
301 gallery = gr.Gallery(
302 label="Generated images", show_label=False, elem_id="gallery"
303 ).style(preview=True, grid=2, object_fit="scale-down")
304 meta_file = gr.File(label="Metadata file")
305 with gr.Column(elem_classes=["container"]):
306 gr.Examples(
307 [
308 ["examples/dog.png", "auto", ""],
309 ["examples/eiffel.jpg", "auto", "tower . lake . grass . sky"],
310 ["examples/eiffel.png", "segment", ""],
311 ["examples/girl.png", "auto", "girl . face"],
312 ["examples/horse.png", "detect", "horse"],
313 ["examples/traffic.jpg", "auto", ""],
314 ],
315 [input_image, task, text_prompt],
316 )
317 gr.HTML(
318 """<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/dragonSwing/annotate-anything?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
319 <p><img src="https://visitor-badge.glitch.me/badge?page_id=dragonswing.annotate-anything" alt="visitors"></p></center>"""
320 )
321
322 run_button.click(
323 fn=process,
324 inputs=[
325 input_image,
326 task,
327 text_prompt,
328 box_threshold,
329 text_threshold,
330 iou_threshold,
331 kernel_size,
332 expand_mask,
333 ],
334 outputs=[gallery, meta_file],
335 )
336
337demo.queue(concurrency_count=2).launch()
338