Update app.py

app.py

@@ -1,65 +1,27 @@
 import gradio as gr
 
 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
 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 huggingface_hub import hf_hub_download
-
-
-# Use this command for evaluate the GLIP-T model
-config_file = "groundingdino/config/GroundingDINO_SwinB_cfg.py"
-ckpt_repo_id = "ShilongLiu/GroundingDINO"
-ckpt_filenmae = "groundingdino_swinb_cogcoor.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=device)
-    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
+# Replace custom imports with Transformers
+from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
+# Add supervision for better visualization
+import supervision as sv
 
-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 ID for Hugging Face
+model_id = "IDEA-Research/grounding-dino-base"
 
-model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae)
+# Load model and processor using Transformers
+device = "cuda" if torch.cuda.is_available() else "cpu"
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device)
 
 def run_grounding(input_image, grounding_caption, box_threshold, text_threshold):
     # Convert numpy array to PIL Image if needed
@@ -69,16 +31,86 @@ def run_grounding(input_image, grounding_caption, box_threshold, text_threshold)
         input_image = Image.fromarray(input_image)
     
     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))
-
+    
+    # Process input using transformers
+    inputs = processor(images=init_image, text=grounding_caption, return_tensors="pt").to(device)
+    
+    # Run inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Post-process results
+    results = processor.post_process_grounded_object_detection(
+        outputs,
+        inputs.input_ids,
+        box_threshold=box_threshold,
+        text_threshold=text_threshold,
+        target_sizes=[init_image.size[::-1]]
+    )
+    
+    result = results[0]
+    
+    # Convert image for supervision visualization
+    image_np = np.array(init_image)
+    
+    # Create detections for supervision
+    boxes = []
+    labels = []
+    confidences = []
+    class_ids = []
+    
+    for i, (box, score, label) in enumerate(zip(result["boxes"], result["scores"], result["labels"])):
+        # Convert box to xyxy format
+        xyxy = box.tolist()
+        boxes.append(xyxy)
+        labels.append(label)
+        confidences.append(float(score))
+        class_ids.append(i)  # Use index as class_id (integer)
+    
+    # Create Detections object for supervision
+    if boxes:
+        detections = sv.Detections(
+            xyxy=np.array(boxes),
+            confidence=np.array(confidences),
+            class_id=np.array(class_ids, dtype=np.int32),  # Ensure it's an integer array
+        )
+        
+        text_scale = sv.calculate_optimal_text_scale(resolution_wh=init_image.size)
+        line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=init_image.size)
+    
+        # Create annotators
+        box_annotator = sv.BoxAnnotator(
+            thickness=2,
+            color=sv.ColorPalette.DEFAULT,
+        )
+        
+        label_annotator = sv.LabelAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            text_color=sv.Color.WHITE,
+            text_scale=text_scale,
+            text_thickness=line_thickness,
+            text_padding=3
+        )
+        
+        # Create formatted labels for each detection
+        formatted_labels = [
+            f"{label}: {conf:.2f}" 
+            for label, conf in zip(labels, confidences)
+        ]
+        
+        # Apply annotations to the image
+        annotated_image = box_annotator.annotate(scene=image_np, detections=detections)
+        annotated_image = label_annotator.annotate(
+            scene=annotated_image, 
+            detections=detections, 
+            labels=formatted_labels
+        )
+    else:
+        annotated_image = image_np
+    
+    # Convert back to PIL Image
+    image_with_box = Image.fromarray(annotated_image)
+    
     return image_with_box
 
 if __name__ == "__main__":
@@ -98,17 +130,16 @@ if __name__ == "__main__":
     with gr.Blocks(css=css) as demo:
         gr.Markdown("<h1><center>Grounding DINO<h1><center>")
         gr.Markdown("<h3><center>Open-World Detection with <a href='https://github.com/IDEA-Research/GroundingDINO'>Grounding DINO</a><h3><center>")
-        gr.Markdown("<h3><center>Running on CPU, so it may take a while to run the model.<h3><center>")
 
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label="Input Image", type="pil")
-                grounding_caption = gr.Textbox(label="Detection Prompt")
+                grounding_caption = gr.Textbox(label="Detection Prompt(VERY important: text queries need to be lowercased + end with a dot, example: a cat. a remote control.)", value="a person. a car.")
                 run_button = gr.Button("Run")
                 
                 with gr.Accordion("Advanced options", open=False):
                     box_threshold = gr.Slider(
-                        minimum=0.0, maximum=1.0, value=0.25, step=0.001, 
+                        minimum=0.0, maximum=1.0, value=0.3, step=0.001, 
                         label="Box Threshold"
                     )
                     text_threshold = gr.Slider(
@@ -129,11 +160,14 @@ if __name__ == "__main__":
         )
         
         gr.Examples(
-            examples=[["this_is_fine.png", "coffee cup", 0.25, 0.25]],
+            examples=[
+                ["000000039769.jpg", "a cat. a remote control.", 0.3, 0.25],
+                ["KakaoTalk_20250430_163200504.jpg", "cup. screen. hand.", 0.3, 0.25]
+                ],
             inputs=[input_image, grounding_caption, box_threshold, text_threshold],
             outputs=[gallery],
             fn=run_grounding,
             cache_examples=True,
         )
     
-    demo.launch(share=args.share, debug=args.debug, show_error=True)
+    demo.launch(share=args.share, debug=args.debug, show_error=True)