add new models

app.py

@@ -10,10 +10,43 @@ def extract_model_short_name(model_id):
 
 model_llmdet_id = "iSEE-Laboratory/llmdet_tiny"
 model_mm_grounding_id = "rziga/mm_grounding_dino_tiny_o365v1_goldg"
+model_omdet_id = "omlab/omdet-turbo-swin-tiny-hf"
+model_owlv2_id = "google/owlv2-large-patch14-ensemble"
 
 model_llmdet_name = extract_model_short_name(model_llmdet_id)
 model_mm_grounding_name = extract_model_short_name(model_mm_grounding_id)
+model_omdet_name = extract_model_short_name(model_omdet_id)
+model_owlv2_name = extract_model_short_name(model_owlv2_id)
 
+@spaces.GPU
+def detect_omdet(image: Image.Image, prompts: list, threshold: float):
+    t0 = time.perf_counter()
+    model_id = model_omdet_id
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processor = AutoProcessor.from_pretrained(model_id)
+    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
+    texts = [prompts]
+    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    results = processor.post_process_grounded_object_detection(
+        outputs,
+        threshold=threshold,
+        target_sizes=[image.size[::-1]]
+    )   
+    result = results[0]
+    annotations = []
+    raw_results = []
+    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
+        if score >= threshold:
+            label_name = prompts[label]
+            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
+            annotations.append(((xmin, ymin, xmax, ymax), f"{label_name} {score:.2f}"))
+            raw_results.append(f"Detected {label_name} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
+    elapsed_ms = (time.perf_counter() - t0) * 1000
+    time_taken = f"**Inference time ({model_omdet_name}):** {elapsed_ms:.0f} ms"
+    raw_text = "\n".join(raw_results) if raw_results else "No detections"
+    return annotations, raw_text, time_taken
 @spaces.GPU
 def detect_llmdet(image: Image.Image, prompts: list, threshold: float):
     t0 = time.perf_counter()
@@ -42,7 +75,6 @@ def detect_llmdet(image: Image.Image, prompts: list, threshold: float):
     time_taken = f"**Inference time ({model_llmdet_name}):** {elapsed_ms:.0f} ms"
     raw_text = "\n".join(raw_results) if raw_results else "No detections"
     return annotations, raw_text, time_taken
-
 @spaces.GPU
 def detect_mm_grounding(image: Image.Image, prompts: list, threshold: float):
     t0 = time.perf_counter()
@@ -72,14 +104,46 @@ def detect_mm_grounding(image: Image.Image, prompts: list, threshold: float):
     raw_text = "\n".join(raw_results) if raw_results else "No detections"
     return annotations, raw_text, time_taken
 
+@spaces.GPU
+def detect_owlv2(image: Image.Image, prompts: list, threshold: float):
+    t0 = time.perf_counter()
+    model_id = model_owlv2_id
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processor = AutoProcessor.from_pretrained(model_id)
+    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
+    texts = [prompts]
+    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    results = processor.post_process_grounded_object_detection(
+        outputs,
+        threshold=threshold,
+        target_sizes=[image.size[::-1]]
+    )
+    result = results[0]
+    annotations = []
+    raw_results = []
+    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
+        if score >= threshold:
+            label_name = prompts[label]
+            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
+            annotations.append(((xmin, ymin, xmax, ymax), f"{label_name} {score:.2f}"))
+            raw_results.append(f"Detected {label_name} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
+    elapsed_ms = (time.perf_counter() - t0) * 1000
+    time_taken = f"**Inference time ({model_owlv2_name}):** {elapsed_ms:.0f} ms"
+    raw_text = "\n".join(raw_results) if raw_results else "No detections"
+    return annotations, raw_text, time_taken
+
 
-def run_detection(image, prompts_str, threshold):
+def run_detection(image, prompts_str, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet):
     if image is None:
         return (None, []), "No detections", "", (None, []), "No detections", ""
     prompts = [p.strip() for p in prompts_str.split(",")]
-    ann_llm, raw_llm, time_llm = detect_llmdet(image, prompts, threshold)
-    ann_mm, raw_mm, time_mm = detect_mm_grounding(image, prompts, threshold)
-    return (image, ann_llm), raw_llm, time_llm, (image, ann_mm), raw_mm, time_mm
+    ann_llm, raw_llm, time_llm = detect_llmdet(image, prompts, threshold_llm)
+    ann_mm, raw_mm, time_mm = detect_mm_grounding(image, prompts, threshold_mm)
+    ann_owlv2, raw_owlv2, time_owlv2 = detect_owlv2(image, prompts, threshold_owlv2)
+    ann_omdet, raw_omdet, time_omdet = detect_omdet(image, prompts, threshold_omdet)
+    return (image, ann_llm), raw_llm, time_llm, (image, ann_mm), raw_mm, time_mm, (image, ann_owlv2), raw_owlv2, time_owlv2, (image, ann_omdet), raw_omdet, time_omdet
 
 with gr.Blocks() as app:
     gr.Markdown("# Zero-Shot Object Detection Arena")
@@ -88,35 +152,52 @@ with gr.Blocks() as app:
         with gr.Column(scale=1):
             image = gr.Image(type="pil", label="Upload an image", height=400)
             prompts = gr.Textbox(label="Prompts (comma-separated)", value="a cat, a remote control")
-            threshold = gr.Slider(label="Confidence Threshold", minimum=0.0, maximum=1.0, step=0.05, value=0.30)
+            with gr.Accordion("Per-model confidence thresholds", open=True):
+                threshold_llm = gr.Slider(label="Threshold for LLMDet", minimum=0.0, maximum=1.0, value=0.3)
+                threshold_mm = gr.Slider(label="Threshold for MM GroundingDINO Tiny", minimum=0.0, maximum=1.0, value=0.3)
+                threshold_owlv2 = gr.Slider(label="Threshold for OwlV2 Large", minimum=0.0, maximum=1.0, value=0.1)
+                threshold_omdet = gr.Slider(label="Threshold for OMDet Turbo Swin Tiny", minimum=0.0, maximum=1.0, value=0.2)
             generate_btn = gr.Button(value="Detect")
-        with gr.Column(scale=2):
-            output_image_llm = gr.AnnotatedImage(label=f"Annotated image for {model_llmdet_name}", height=400)
-            output_text_llm = gr.Textbox(label=f"Model detections for {model_llmdet_name}", lines=10)
-            output_time_llm = gr.Markdown()
-        with gr.Column(scale=2):
-            output_image_mm = gr.AnnotatedImage(label=f"Annotated image for {model_mm_grounding_name}", height=400)
-            output_text_mm = gr.Textbox(label=f"Model detections for {model_mm_grounding_name}", lines=10)
-            output_time_mm = gr.Markdown()
+        with gr.Row():
+            with gr.Column(scale=2):
+                output_image_llm = gr.AnnotatedImage(label=f"Annotated image for {model_llmdet_name}", height=400)
+                output_text_llm = gr.Textbox(label=f"Model detections for {model_llmdet_name}", lines=5)
+                output_time_llm = gr.Markdown()
+            with gr.Column(scale=2):
+                output_image_mm = gr.AnnotatedImage(label=f"Annotated image for {model_mm_grounding_name}", height=400)
+                output_text_mm = gr.Textbox(label=f"Model detections for {model_mm_grounding_name}", lines=5)
+                output_time_mm = gr.Markdown()
+        with gr.Row():
+            with gr.Column(scale=2):
+                output_image_owlv2 = gr.AnnotatedImage(label=f"Annotated image for {model_owlv2_name}", height=400)
+                output_text_owlv2 = gr.Textbox(label=f"Model detections for {model_owlv2_name}", lines=5)
+                output_time_owlv2 = gr.Markdown()
+            with gr.Column(scale=2):
+                output_image_omdet = gr.AnnotatedImage(label=f"Annotated image for {model_omdet_name}", height=400)
+                output_text_omdet = gr.Textbox(label=f"Model detections for {model_omdet_name}", lines=5)
+                output_time_omdet = gr.Markdown()
     gr.Markdown("### Examples")
     example_data = [
-        ["http://images.cocodataset.org/val2017/000000039769.jpg", "a cat, a remote control", 0.4],
-        ["http://images.cocodataset.org/val2017/000000000139.jpg", "a person, a tv, a remote", 0.3],
+        ["http://images.cocodataset.org/val2017/000000039769.jpg", "a cat, a remote control", 0.30, 0.30, 0.10, 0.30],
+        ["http://images.cocodataset.org/val2017/000000000139.jpg", "a person, a tv, a remote", 0.35, 0.30, 0.12, 0.30],
     ]
+
     gr.Examples(
         examples=example_data,
-        inputs=[image, prompts, threshold],
-        label="Click an example to populate the input",
+        inputs=[image, prompts, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet],
+        label="Click an example to populate the inputs",
     )
+    inputs = [image, prompts, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet]
+    outputs = [output_image_llm, output_text_llm, output_time_llm, output_image_mm, output_text_mm, output_time_mm, output_image_owlv2, output_text_owlv2, output_time_owlv2, output_image_omdet, output_text_omdet, output_time_omdet]
     generate_btn.click(
         fn=run_detection,
-        inputs=[image, prompts, threshold],
-        outputs=[output_image_llm, output_text_llm, output_time_llm, output_image_mm, output_text_mm, output_time_mm],
+        inputs=inputs,
+        outputs=outputs,
     )
     image.upload(
         fn=run_detection,
-        inputs=[image, prompts, threshold],
-        outputs=[output_image_llm, output_text_llm, output_time_llm, output_image_mm, output_text_mm, output_time_mm],
+        inputs=inputs,
+        outputs=outputs,
     )
 
 app.launch()

requirements.txt

@@ -4,4 +4,6 @@ pillow
 spaces
 gradio
 transformers
-accelerate
+accelerate
+scipy
+timm