Compare commits
3 Commits
main
...
release_co
| Author | SHA1 | Date |
|---|---|---|
SlongLiu
|
5a826116c8 | 2 years ago |
|
SlongLiu
|
08bd3c2e9b | 2 years ago |
|
SlongLiu
|
cc02608066 | 2 years ago |
15 changed files with 72 additions and 1863 deletions
@ -1,125 +0,0 @@ |
|||||||
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 warnings |
|
||||||
|
|
||||||
import torch |
|
||||||
|
|
||||||
# prepare the environment |
|
||||||
os.system("python setup.py build develop --user") |
|
||||||
os.system("pip install packaging==21.3") |
|
||||||
os.system("pip install gradio") |
|
||||||
|
|
||||||
|
|
||||||
warnings.filterwarnings("ignore") |
|
||||||
|
|
||||||
import gradio as gr |
|
||||||
|
|
||||||
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 huggingface_hub import hf_hub_download |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# Use this command for evaluate the Grounding DINO model |
|
||||||
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) # 3, h, w |
|
||||||
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) # 3, h, w |
|
||||||
return image |
|
||||||
|
|
||||||
model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae) |
|
||||||
|
|
||||||
def run_grounding(input_image, grounding_caption, box_threshold, text_threshold): |
|
||||||
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) |
|
||||||
|
|
||||||
# run grounidng |
|
||||||
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() |
|
||||||
|
|
||||||
block = gr.Blocks().queue() |
|
||||||
with block: |
|
||||||
gr.Markdown("# [Grounding DINO](https://github.com/IDEA-Research/GroundingDINO)") |
|
||||||
gr.Markdown("### Open-World Detection with Grounding DINO") |
|
||||||
|
|
||||||
with gr.Row(): |
|
||||||
with gr.Column(): |
|
||||||
input_image = gr.Image(source='upload', type="pil") |
|
||||||
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", |
|
||||||
# label="grounding results" |
|
||||||
).style(full_width=True, full_height=True) |
|
||||||
# gallery = gr.Gallery(label="Generated images", show_label=False).style( |
|
||||||
# grid=[1], height="auto", container=True, full_width=True, full_height=True) |
|
||||||
|
|
||||||
run_button.click(fn=run_grounding, inputs=[ |
|
||||||
input_image, grounding_caption, box_threshold, text_threshold], outputs=[gallery]) |
|
||||||
|
|
||||||
|
|
||||||
block.launch(server_name='0.0.0.0', server_port=7579, debug=args.debug, share=args.share) |
|
||||||
|
|
||||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@ -1,43 +0,0 @@ |
|||||||
batch_size = 1 |
|
||||||
modelname = "groundingdino" |
|
||||||
backbone = "swin_B_384_22k" |
|
||||||
position_embedding = "sine" |
|
||||||
pe_temperatureH = 20 |
|
||||||
pe_temperatureW = 20 |
|
||||||
return_interm_indices = [1, 2, 3] |
|
||||||
backbone_freeze_keywords = None |
|
||||||
enc_layers = 6 |
|
||||||
dec_layers = 6 |
|
||||||
pre_norm = False |
|
||||||
dim_feedforward = 2048 |
|
||||||
hidden_dim = 256 |
|
||||||
dropout = 0.0 |
|
||||||
nheads = 8 |
|
||||||
num_queries = 900 |
|
||||||
query_dim = 4 |
|
||||||
num_patterns = 0 |
|
||||||
num_feature_levels = 4 |
|
||||||
enc_n_points = 4 |
|
||||||
dec_n_points = 4 |
|
||||||
two_stage_type = "standard" |
|
||||||
two_stage_bbox_embed_share = False |
|
||||||
two_stage_class_embed_share = False |
|
||||||
transformer_activation = "relu" |
|
||||||
dec_pred_bbox_embed_share = True |
|
||||||
dn_box_noise_scale = 1.0 |
|
||||||
dn_label_noise_ratio = 0.5 |
|
||||||
dn_label_coef = 1.0 |
|
||||||
dn_bbox_coef = 1.0 |
|
||||||
embed_init_tgt = True |
|
||||||
dn_labelbook_size = 2000 |
|
||||||
max_text_len = 256 |
|
||||||
text_encoder_type = "bert-base-uncased" |
|
||||||
use_text_enhancer = True |
|
||||||
use_fusion_layer = True |
|
||||||
use_checkpoint = True |
|
||||||
use_transformer_ckpt = True |
|
||||||
use_text_cross_attention = True |
|
||||||
text_dropout = 0.0 |
|
||||||
fusion_dropout = 0.0 |
|
||||||
fusion_droppath = 0.1 |
|
||||||
sub_sentence_present = True |
|
||||||
@ -1,242 +0,0 @@ |
|||||||
from typing import Tuple, List |
|
||||||
|
|
||||||
import cv2 |
|
||||||
import numpy as np |
|
||||||
import supervision as sv |
|
||||||
import torch |
|
||||||
from PIL import Image |
|
||||||
from torchvision.ops import box_convert |
|
||||||
|
|
||||||
import groundingdino.datasets.transforms as T |
|
||||||
from groundingdino.models import build_model |
|
||||||
from groundingdino.util.misc import clean_state_dict |
|
||||||
from groundingdino.util.slconfig import SLConfig |
|
||||||
from groundingdino.util.utils import get_phrases_from_posmap |
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------------------------------------------------- |
|
||||||
# OLD API |
|
||||||
# ---------------------------------------------------------------------------------------------------------------------- |
|
||||||
|
|
||||||
|
|
||||||
def preprocess_caption(caption: str) -> str: |
|
||||||
result = caption.lower().strip() |
|
||||||
if result.endswith("."): |
|
||||||
return result |
|
||||||
return result + "." |
|
||||||
|
|
||||||
|
|
||||||
def load_model(model_config_path: str, model_checkpoint_path: str, device: str = "cuda"): |
|
||||||
args = SLConfig.fromfile(model_config_path) |
|
||||||
args.device = device |
|
||||||
model = build_model(args) |
|
||||||
checkpoint = torch.load(model_checkpoint_path, map_location="cpu") |
|
||||||
model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) |
|
||||||
model.eval() |
|
||||||
return model |
|
||||||
|
|
||||||
|
|
||||||
def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]: |
|
||||||
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_source = Image.open(image_path).convert("RGB") |
|
||||||
image = np.asarray(image_source) |
|
||||||
image_transformed, _ = transform(image_source, None) |
|
||||||
return image, image_transformed |
|
||||||
|
|
||||||
|
|
||||||
def predict( |
|
||||||
model, |
|
||||||
image: torch.Tensor, |
|
||||||
caption: str, |
|
||||||
box_threshold: float, |
|
||||||
text_threshold: float, |
|
||||||
device: str = "cuda" |
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor, List[str]]: |
|
||||||
caption = preprocess_caption(caption=caption) |
|
||||||
|
|
||||||
model = model.to(device) |
|
||||||
image = image.to(device) |
|
||||||
|
|
||||||
with torch.no_grad(): |
|
||||||
outputs = model(image[None], captions=[caption]) |
|
||||||
|
|
||||||
prediction_logits = outputs["pred_logits"].cpu().sigmoid()[0] # prediction_logits.shape = (nq, 256) |
|
||||||
prediction_boxes = outputs["pred_boxes"].cpu()[0] # prediction_boxes.shape = (nq, 4) |
|
||||||
|
|
||||||
mask = prediction_logits.max(dim=1)[0] > box_threshold |
|
||||||
logits = prediction_logits[mask] # logits.shape = (n, 256) |
|
||||||
boxes = prediction_boxes[mask] # boxes.shape = (n, 4) |
|
||||||
|
|
||||||
tokenizer = model.tokenizer |
|
||||||
tokenized = tokenizer(caption) |
|
||||||
|
|
||||||
phrases = [ |
|
||||||
get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer).replace('.', '') |
|
||||||
for logit |
|
||||||
in logits |
|
||||||
] |
|
||||||
|
|
||||||
return boxes, logits.max(dim=1)[0], phrases |
|
||||||
|
|
||||||
|
|
||||||
def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str]) -> np.ndarray: |
|
||||||
h, w, _ = image_source.shape |
|
||||||
boxes = boxes * torch.Tensor([w, h, w, h]) |
|
||||||
xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() |
|
||||||
detections = sv.Detections(xyxy=xyxy) |
|
||||||
|
|
||||||
labels = [ |
|
||||||
f"{phrase} {logit:.2f}" |
|
||||||
for phrase, logit |
|
||||||
in zip(phrases, logits) |
|
||||||
] |
|
||||||
|
|
||||||
box_annotator = sv.BoxAnnotator() |
|
||||||
annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR) |
|
||||||
annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels) |
|
||||||
return annotated_frame |
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------------------------------------------------- |
|
||||||
# NEW API |
|
||||||
# ---------------------------------------------------------------------------------------------------------------------- |
|
||||||
|
|
||||||
|
|
||||||
class Model: |
|
||||||
|
|
||||||
def __init__( |
|
||||||
self, |
|
||||||
model_config_path: str, |
|
||||||
model_checkpoint_path: str, |
|
||||||
device: str = "cuda" |
|
||||||
): |
|
||||||
self.model = load_model( |
|
||||||
model_config_path=model_config_path, |
|
||||||
model_checkpoint_path=model_checkpoint_path, |
|
||||||
device=device |
|
||||||
).to(device) |
|
||||||
self.device = device |
|
||||||
|
|
||||||
def predict_with_caption( |
|
||||||
self, |
|
||||||
image: np.ndarray, |
|
||||||
caption: str, |
|
||||||
box_threshold: float = 0.35, |
|
||||||
text_threshold: float = 0.25 |
|
||||||
) -> Tuple[sv.Detections, List[str]]: |
|
||||||
""" |
|
||||||
import cv2 |
|
||||||
|
|
||||||
image = cv2.imread(IMAGE_PATH) |
|
||||||
|
|
||||||
model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) |
|
||||||
detections, labels = model.predict_with_caption( |
|
||||||
image=image, |
|
||||||
caption=caption, |
|
||||||
box_threshold=BOX_THRESHOLD, |
|
||||||
text_threshold=TEXT_THRESHOLD |
|
||||||
) |
|
||||||
|
|
||||||
import supervision as sv |
|
||||||
|
|
||||||
box_annotator = sv.BoxAnnotator() |
|
||||||
annotated_image = box_annotator.annotate(scene=image, detections=detections, labels=labels) |
|
||||||
""" |
|
||||||
processed_image = Model.preprocess_image(image_bgr=image).to(self.device) |
|
||||||
boxes, logits, phrases = predict( |
|
||||||
model=self.model, |
|
||||||
image=processed_image, |
|
||||||
caption=caption, |
|
||||||
box_threshold=box_threshold, |
|
||||||
text_threshold=text_threshold) |
|
||||||
source_h, source_w, _ = image.shape |
|
||||||
detections = Model.post_process_result( |
|
||||||
source_h=source_h, |
|
||||||
source_w=source_w, |
|
||||||
boxes=boxes, |
|
||||||
logits=logits) |
|
||||||
return detections, phrases |
|
||||||
|
|
||||||
def predict_with_classes( |
|
||||||
self, |
|
||||||
image: np.ndarray, |
|
||||||
classes: List[str], |
|
||||||
box_threshold: float, |
|
||||||
text_threshold: float |
|
||||||
) -> sv.Detections: |
|
||||||
""" |
|
||||||
import cv2 |
|
||||||
|
|
||||||
image = cv2.imread(IMAGE_PATH) |
|
||||||
|
|
||||||
model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) |
|
||||||
detections = model.predict_with_classes( |
|
||||||
image=image, |
|
||||||
classes=CLASSES, |
|
||||||
box_threshold=BOX_THRESHOLD, |
|
||||||
text_threshold=TEXT_THRESHOLD |
|
||||||
) |
|
||||||
|
|
||||||
|
|
||||||
import supervision as sv |
|
||||||
|
|
||||||
box_annotator = sv.BoxAnnotator() |
|
||||||
annotated_image = box_annotator.annotate(scene=image, detections=detections) |
|
||||||
""" |
|
||||||
caption = ", ".join(classes) |
|
||||||
processed_image = Model.preprocess_image(image_bgr=image).to(self.device) |
|
||||||
boxes, logits, phrases = predict( |
|
||||||
model=self.model, |
|
||||||
image=processed_image, |
|
||||||
caption=caption, |
|
||||||
box_threshold=box_threshold, |
|
||||||
text_threshold=text_threshold) |
|
||||||
source_h, source_w, _ = image.shape |
|
||||||
detections = Model.post_process_result( |
|
||||||
source_h=source_h, |
|
||||||
source_w=source_w, |
|
||||||
boxes=boxes, |
|
||||||
logits=logits) |
|
||||||
class_id = Model.phrases2classes(phrases=phrases, classes=classes) |
|
||||||
detections.class_id = class_id |
|
||||||
return detections |
|
||||||
|
|
||||||
@staticmethod |
|
||||||
def preprocess_image(image_bgr: np.ndarray) -> torch.Tensor: |
|
||||||
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_pillow = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)) |
|
||||||
image_transformed, _ = transform(image_pillow, None) |
|
||||||
return image_transformed |
|
||||||
|
|
||||||
@staticmethod |
|
||||||
def post_process_result( |
|
||||||
source_h: int, |
|
||||||
source_w: int, |
|
||||||
boxes: torch.Tensor, |
|
||||||
logits: torch.Tensor |
|
||||||
) -> sv.Detections: |
|
||||||
boxes = boxes * torch.Tensor([source_w, source_h, source_w, source_h]) |
|
||||||
xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() |
|
||||||
confidence = logits.numpy() |
|
||||||
return sv.Detections(xyxy=xyxy, confidence=confidence) |
|
||||||
|
|
||||||
@staticmethod |
|
||||||
def phrases2classes(phrases: List[str], classes: List[str]) -> np.ndarray: |
|
||||||
class_ids = [] |
|
||||||
for phrase in phrases: |
|
||||||
try: |
|
||||||
class_ids.append(classes.index(phrase)) |
|
||||||
except ValueError: |
|
||||||
class_ids.append(None) |
|
||||||
return np.array(class_ids) |
|
||||||
@ -1,10 +1 @@ |
|||||||
torch |
transformers==4.5.1 |
||||||
torchvision |
|
||||||
transformers |
|
||||||
addict |
|
||||||
yapf |
|
||||||
timm |
|
||||||
numpy |
|
||||||
opencv-python |
|
||||||
supervision==0.4.0 |
|
||||||
pycocotools |
|
||||||
Loading…
Reference in new issue