|
import argparse |
|
from functools import partial |
|
import cv2 |
|
import requests |
|
import os |
|
from io import BytesIO |
|
from PIL import Image |
|
import numpy as np |
|
from pathlib import Path |
|
import gradio as gr |
|
|
|
import warnings |
|
|
|
import torch |
|
|
|
os.system("python setup.py build develop --user") |
|
os.system("pip install packaging==21.3") |
|
warnings.filterwarnings("ignore") |
|
|
|
|
|
from groundingdino.models import build_model |
|
from groundingdino.util.slconfig import SLConfig |
|
from groundingdino.util.utils import clean_state_dict |
|
from groundingdino.util.inference import annotate, load_image, predict |
|
import groundingdino.datasets.transforms as T |
|
from groundingdino.util import box_ops |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
|
|
|
|
config_file = "groundingdino/config/GroundingDINO_SwinT_OGC.py" |
|
ckpt_repo_id = "ShilongLiu/GroundingDINO" |
|
ckpt_filenmae = "groundingdino_swint_ogc.pth" |
|
|
|
|
|
def load_model_hf(model_config_path, repo_id, filename, device='cpu'): |
|
args = SLConfig.fromfile(model_config_path) |
|
model = build_model(args) |
|
args.device = device |
|
|
|
cache_file = hf_hub_download(repo_id=repo_id, filename=filename) |
|
checkpoint = torch.load(cache_file, map_location='cpu') |
|
log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False) |
|
print("Model loaded from {} \n => {}".format(cache_file, log)) |
|
_ = model.eval() |
|
return model |
|
|
|
def image_transform_grounding(init_image): |
|
transform = T.Compose([ |
|
T.RandomResize([800], max_size=1333), |
|
T.ToTensor(), |
|
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
|
]) |
|
image, _ = transform(init_image, None) |
|
return init_image, image |
|
|
|
def image_transform_grounding_for_vis(init_image): |
|
transform = T.Compose([ |
|
T.RandomResize([800], max_size=1333), |
|
]) |
|
image, _ = transform(init_image, None) |
|
return image |
|
|
|
model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae) |
|
|
|
def segment(image, sam_model, boxes): |
|
H, W, _ = image.shape |
|
boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H]) |
|
|
|
transformed_boxes = sam_model.transform.apply_boxes_torch(boxes_xyxy.to(device), image.shape[:2]) |
|
masks, _, _ = sam_model.predict_torch( |
|
point_coords = None, |
|
point_labels = None, |
|
boxes = transformed_boxes, |
|
multimask_output = False, |
|
) |
|
return masks.cpu() |
|
|
|
|
|
def draw_mask(mask, image, random_color=True): |
|
if random_color: |
|
color = np.concatenate([np.random.random(3), np.array([0.8])], axis=0) |
|
else: |
|
color = np.array([30/255, 144/255, 255/255, 0.6]) |
|
h, w = mask.shape[-2:] |
|
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) |
|
|
|
annotated_frame_pil = Image.fromarray(image).convert("RGBA") |
|
mask_image_pil = Image.fromarray((mask_image.cpu().numpy() * 255).astype(np.uint8)).convert("RGBA") |
|
|
|
return np.array(Image.alpha_composite(annotated_frame_pil, mask_image_pil)) |
|
|
|
|
|
def run_grounding(input_image,choice, grounding_caption, box_threshold, text_threshold,do_segmentation): |
|
init_image = input_image.convert("RGB") |
|
original_size = init_image.size |
|
|
|
_, image_tensor = image_transform_grounding(init_image) |
|
image_pil: Image = image_transform_grounding_for_vis(init_image) |
|
|
|
if choice == 'segment': |
|
boxes, logits, phrases = predict(model, image_tensor, grounding_caption, box_threshold, text_threshold, device='cpu') |
|
segmented_frame_masks = segment(image_tensor, model, boxes=boxes) |
|
annotated_frame_with_mask = draw_mask(segmented_frame_masks[0][0], annotated_frame) |
|
else: |
|
|
|
boxes, logits, phrases = predict(model, image_tensor, grounding_caption, box_threshold, text_threshold, device='cpu') |
|
annotated_frame = annotate(image_source=np.asarray(image_pil), boxes=boxes, logits=logits, phrases=phrases) |
|
|
|
image_with_box = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)) |
|
|
|
return image_with_box |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
parser = argparse.ArgumentParser("Grounding DINO demo", add_help=True) |
|
parser.add_argument("--debug", action="store_true", help="using debug mode") |
|
parser.add_argument("--share", action="store_true", help="share the app") |
|
args = parser.parse_args() |
|
css = """ |
|
#mkd { |
|
height: 500px; |
|
overflow: auto; |
|
border: 1px solid #ccc; |
|
} |
|
""" |
|
block = gr.Blocks(css=css).queue() |
|
with block: |
|
gr.Markdown("<h1><center>Grounding DINO<h1><center>") |
|
gr.Markdown("<h3><center>Open-World Detection with <a href='https://github.com/Arulkumar03/SOTA-Grounding-DINO.ipynb'>Grounding DINO</a><h3><center>") |
|
gr.Markdown("<h3><center>Note the model runs on CPU, so it may take a while to run the model.<h3><center>") |
|
|
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image(source='upload', type="pil") |
|
choice = gr.Radio( |
|
["segment", "classify"], default="segment", label="Choose Operation" |
|
) |
|
grounding_caption = gr.Textbox(label="Detection Prompt") |
|
run_button = gr.Button(label="Run") |
|
with gr.Accordion("Advanced options", open=False): |
|
box_threshold = gr.Slider( |
|
label="Box Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.001 |
|
) |
|
text_threshold = gr.Slider( |
|
label="Text Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.001 |
|
) |
|
|
|
with gr.Column(): |
|
gallery = gr.outputs.Image( |
|
type="pil", |
|
|
|
).style(full_width=True, full_height=True) |
|
|
|
run_button.click(fn=run_grounding, inputs=[ |
|
input_image, choice, grounding_caption, box_threshold, text_threshold], outputs=[gallery]) |
|
gr.Examples( |
|
[["watermelon.jpg", "segment", "watermelon", 0.25, 0.25]], |
|
inputs=[input_image, choice, grounding_caption, box_threshold, text_threshold], |
|
outputs=[gallery], |
|
fn=run_grounding, |
|
cache_examples=True, |
|
label='Try this example input!' |
|
) |
|
block.launch(share=False, show_api=False, show_error=True) |
|
|