Spaces:
Running
Running
| import functools | |
| import json | |
| import os | |
| import sys | |
| import tempfile | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import supervision as sv | |
| import torch | |
| from PIL import Image | |
| from segment_anything import build_sam | |
| from segment_anything import SamAutomaticMaskGenerator | |
| from segment_anything import SamPredictor | |
| from supervision.detection.utils import mask_to_polygons | |
| from supervision.detection.utils import xywh_to_xyxy | |
| if os.environ.get("IS_MY_DEBUG") is None: | |
| os.system("pip install -e GroundingDINO") | |
| sys.path.append("tag2text") | |
| sys.path.append("GroundingDINO") | |
| from groundingdino.util.inference import Model as DinoModel | |
| from tag2text.models import tag2text | |
| from config import * | |
| from utils import download_file_hf, detect, segment, generate_tags | |
| if not os.path.exists(abs_weight_dir): | |
| os.makedirs(abs_weight_dir, exist_ok=True) | |
| sam_checkpoint = os.path.join(abs_weight_dir, sam_dict[default_sam]["checkpoint_file"]) | |
| if not os.path.exists(sam_checkpoint): | |
| os.system(f"wget {sam_dict[default_sam]['checkpoint_url']} -O {sam_checkpoint}") | |
| tag2text_checkpoint = os.path.join( | |
| abs_weight_dir, tag2text_dict[default_tag2text]["checkpoint_file"] | |
| ) | |
| if not os.path.exists(tag2text_checkpoint): | |
| os.system( | |
| f"wget {tag2text_dict[default_tag2text]['checkpoint_url']} -O {tag2text_checkpoint}" | |
| ) | |
| dino_checkpoint = os.path.join( | |
| abs_weight_dir, dino_dict[default_dino]["checkpoint_file"] | |
| ) | |
| dino_config_file = os.path.join(abs_weight_dir, dino_dict[default_dino]["config_file"]) | |
| if not os.path.exists(dino_checkpoint): | |
| dino_repo_id = dino_dict[default_dino]["repo_id"] | |
| download_file_hf( | |
| repo_id=dino_repo_id, | |
| filename=dino_dict[default_dino]["config_file"], | |
| cache_dir=weight_dir, | |
| ) | |
| download_file_hf( | |
| repo_id=dino_repo_id, | |
| filename=dino_dict[default_dino]["checkpoint_file"], | |
| cache_dir=weight_dir, | |
| ) | |
| # load model | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| tag2text_model = tag2text.tag2text_caption( | |
| pretrained=tag2text_checkpoint, | |
| image_size=384, | |
| vit="swin_b", | |
| delete_tag_index=delete_tag_index, | |
| ) | |
| # threshold for tagging | |
| # we reduce the threshold to obtain more tags | |
| tag2text_model.threshold = 0.64 | |
| tag2text_model.to(device) | |
| tag2text_model.eval() | |
| sam = build_sam(checkpoint=sam_checkpoint) | |
| sam.to(device=device) | |
| sam_predictor = SamPredictor(sam) | |
| sam_automask_generator = SamAutomaticMaskGenerator(sam) | |
| grounding_dino_model = DinoModel( | |
| model_config_path=dino_config_file, | |
| model_checkpoint_path=dino_checkpoint, | |
| device=device, | |
| ) | |
| def process( | |
| image_path, | |
| task, | |
| prompt, | |
| box_threshold, | |
| text_threshold, | |
| iou_threshold, | |
| kernel_size, | |
| expand_mask, | |
| ): | |
| global tag2text_model, sam_predictor, sam_automask_generator, grounding_dino_model, device | |
| output_gallery = [] | |
| detections = None | |
| metadata = {"image": {}, "annotations": []} | |
| try: | |
| # Load image | |
| image = Image.open(image_path) | |
| image_pil = image.convert("RGB") | |
| image = np.array(image_pil) | |
| orig_image = image.copy() | |
| # Extract image metadata | |
| filename = os.path.basename(image_path) | |
| h, w = image.shape[:2] | |
| metadata["image"]["file_name"] = filename | |
| metadata["image"]["width"] = w | |
| metadata["image"]["height"] = h | |
| # Generate tags | |
| if task in ["auto", "detection"] and prompt == "": | |
| tags, caption = generate_tags(tag2text_model, image_pil, "None", device) | |
| prompt = " . ".join(tags) | |
| print(f"Caption: {caption}") | |
| print(f"Tags: {tags}") | |
| # ToDo: Extract metadata | |
| metadata["image"]["caption"] = caption | |
| metadata["image"]["tags"] = tags | |
| if prompt: | |
| metadata["prompt"] = prompt | |
| print(f"Prompt: {prompt}") | |
| # Detect boxes | |
| if prompt != "": | |
| detections, phrases, classes = detect( | |
| grounding_dino_model, | |
| image, | |
| caption=prompt, | |
| box_threshold=box_threshold, | |
| text_threshold=text_threshold, | |
| iou_threshold=iou_threshold, | |
| post_process=True, | |
| ) | |
| print(phrases) | |
| # Draw boxes | |
| box_annotator = sv.BoxAnnotator() | |
| labels = [ | |
| f"{phrases[i]} {detections.confidence[i]:0.2f}" | |
| for i in range(len(phrases)) | |
| ] | |
| image = box_annotator.annotate( | |
| scene=image, detections=detections, labels=labels | |
| ) | |
| output_gallery.append(image) | |
| # Segmentation | |
| if task in ["auto", "segment"]: | |
| kernel = cv2.getStructuringElement( | |
| cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1) | |
| ) | |
| if detections: | |
| masks, scores = segment( | |
| sam_predictor, image=orig_image, boxes=detections.xyxy | |
| ) | |
| if expand_mask: | |
| masks = [ | |
| cv2.dilate(mask.astype(np.uint8), kernel) for mask in masks | |
| ] | |
| else: | |
| masks = [ | |
| cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel) | |
| for mask in masks | |
| ] | |
| detections.mask = masks | |
| binary_mask = functools.reduce( | |
| lambda x, y: x + y, detections.mask | |
| ).astype(bool) | |
| else: | |
| masks = sam_automask_generator.generate(orig_image) | |
| sorted_generated_masks = sorted( | |
| masks, key=lambda x: x["area"], reverse=True | |
| ) | |
| xywh = np.array([mask["bbox"] for mask in sorted_generated_masks]) | |
| scores = np.array( | |
| [mask["predicted_iou"] for mask in sorted_generated_masks] | |
| ) | |
| if expand_mask: | |
| mask = np.array( | |
| [ | |
| cv2.dilate(mask["segmentation"].astype(np.uint8), kernel) | |
| for mask in sorted_generated_masks | |
| ] | |
| ) | |
| else: | |
| mask = np.array( | |
| [mask["segmentation"] for mask in sorted_generated_masks] | |
| ) | |
| detections = sv.Detections( | |
| xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask | |
| ) | |
| binary_mask = None | |
| mask_annotator = sv.MaskAnnotator() | |
| mask_image = np.zeros_like(image, dtype=np.uint8) | |
| mask_image = mask_annotator.annotate( | |
| mask_image, detections=detections, opacity=1 | |
| ) | |
| annotated_image = mask_annotator.annotate(image, detections=detections) | |
| output_gallery.append(mask_image) | |
| if binary_mask is not None: | |
| binary_mask_image = binary_mask * 255 | |
| cutout_image = np.expand_dims(binary_mask, axis=-1) * orig_image | |
| output_gallery.append(binary_mask_image) | |
| output_gallery.append(cutout_image) | |
| output_gallery.append(annotated_image) | |
| # ToDo: Extract metadata | |
| if detections: | |
| i = 0 | |
| for (xyxy, mask, confidence, _, _), area, box_area in zip( | |
| detections, detections.area, detections.box_area | |
| ): | |
| annotation = { | |
| "id": i + 1, | |
| "bbox": [int(x) for x in xyxy], | |
| "box_area": float(box_area), | |
| } | |
| if confidence: | |
| annotation["confidence"] = float(confidence) | |
| annotation["label"] = phrases[i] | |
| if mask is not None: | |
| # annotation["segmentation"] = mask_to_polygons(mask) | |
| annotation["area"] = int(area) | |
| annotation["predicted_iou"] = float(scores[i]) | |
| metadata["annotations"].append(annotation) | |
| i += 1 | |
| meta_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json") | |
| meta_file_path = meta_file.name | |
| with open(meta_file_path, "w", encoding="utf-8") as fp: | |
| json.dump(metadata, fp) | |
| return output_gallery, meta_file_path | |
| except Exception as error: | |
| raise gr.Error(f"global exception: {error}") | |
| title = "Annotate Anything" | |
| with gr.Blocks(css="style.css", title=title) as demo: | |
| with gr.Row(elem_classes=["container"]): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(type="filepath", label="Input") | |
| task = gr.Dropdown( | |
| ["detect", "segment", "auto"], value="auto", label="task_type" | |
| ) | |
| text_prompt = gr.Textbox( | |
| label="Detection Prompt", | |
| info="To detect multiple objects, seperating each name with '.', like this: cat . dog . chair ", | |
| ) | |
| with gr.Accordion("Advanced parameters", open=False): | |
| box_threshold = gr.Slider( | |
| minimum=0, | |
| maximum=1, | |
| value=0.3, | |
| step=0.05, | |
| label="Box threshold", | |
| ) | |
| text_threshold = gr.Slider( | |
| minimum=0, | |
| maximum=1, | |
| value=0.25, | |
| step=0.05, | |
| label="Text threshold", | |
| ) | |
| iou_threshold = gr.Slider( | |
| minimum=0, | |
| maximum=1, | |
| value=0.5, | |
| step=0.05, | |
| label="IOU threshold", | |
| info="Intersection over Union threshold", | |
| ) | |
| kernel_size = gr.Slider( | |
| minimum=1, | |
| maximum=5, | |
| value=2, | |
| step=1, | |
| label="Kernel size", | |
| info="Use to smooth segment masks", | |
| ) | |
| expand_mask = gr.Checkbox( | |
| label="Expand mask", | |
| ) | |
| run_button = gr.Button(label="Run") | |
| with gr.Column(scale=2): | |
| gallery = gr.Gallery( | |
| label="Generated images", show_label=False, elem_id="gallery" | |
| ).style(preview=True, grid=2, object_fit="scale-down") | |
| meta_file = gr.File(label="Metadata file") | |
| with gr.Column(elem_classes=["container"]): | |
| gr.Examples( | |
| [ | |
| ["examples/dog.png", "auto", ""], | |
| ["examples/eiffel.jpg", "auto", "tower . lake . grass . sky"], | |
| ["examples/eiffel.png", "segment", ""], | |
| ["examples/girl.png", "auto", "girl . face"], | |
| ["examples/horse.png", "detect", "horse"], | |
| ["examples/traffic.jpg", "auto", ""], | |
| ], | |
| [input_image, task, text_prompt], | |
| ) | |
| gr.HTML( | |
| """<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> | |
| <p><img src="https://visitor-badge.glitch.me/badge?page_id=dragonswing.annotate-anything" alt="visitors"></p></center>""" | |
| ) | |
| run_button.click( | |
| fn=process, | |
| inputs=[ | |
| input_image, | |
| task, | |
| text_prompt, | |
| box_threshold, | |
| text_threshold, | |
| iou_threshold, | |
| kernel_size, | |
| expand_mask, | |
| ], | |
| outputs=[gallery, meta_file], | |
| ) | |
| demo.queue(concurrency_count=2).launch() | |