Add Object Detection tab (full-frame D-FINE, class summary)

Third tab after PPE: run any of the 9 D-FINE models on the whole
image with a threshold slider, draw one colour-coded box per
detection (class + score), and list per-class counts below.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -1,10 +1,10 @@
-""" Gradio app: D-FINE + SigLIP Classify, plus PPE Detection. """
+""" Gradio app: D-FINE + SigLIP Classify, PPE Detection, and D-FINE Object Detection. """
 
 import os
 import gradio as gr
 from pathlib import Path
 
-from dfine_jina_pipeline import run_single_image
+from dfine_jina_pipeline import run_single_image, run_object_detection
 from ppe_compliance import run_ppe_compliance
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -43,6 +43,21 @@ def run_dfine_classify(image, dfine_threshold, dfine_model_choice, classifier_ch
     return [(g, None) for g in (group_crops or [])], [(k, None) for k in (known_crops or [])], status or ""
 
 
+def run_objdet(image, dfine_model_choice, det_threshold):
+    """Run a chosen D-FINE model on the full image.
+    Returns (annotated_image, detected-classes summary).
+    """
+    if image is None:
+        return None, "Upload an image."
+
+    dfine_model = dfine_model_choice.strip().lower() if dfine_model_choice else "medium-obj2coco"
+    return run_object_detection(
+        image,
+        dfine_model=dfine_model,
+        det_threshold=float(det_threshold),
+    )
+
+
 IMG_HEIGHT = 400
 
 
@@ -226,6 +241,67 @@ with gr.Blocks(title="Small Object Detection") as app:
             concurrency_limit=1,
         )
 
+      with gr.Tab("Object Detection"):
+
+        gr.Markdown(
+            "Run a **D-FINE** model on the whole image and see every detection. "
+            "Pick a D-FINE model and detection threshold; each object is drawn with a "
+            "colour-coded box (class + score) and the classes found are listed below."
+        )
+
+        with gr.Row():
+
+            with gr.Column(scale=1):
+
+                inp_objdet = gr.Image(
+                    type="pil",
+                    label="Input image",
+                    height=IMG_HEIGHT,
+                )
+
+                objdet_model_dropdown = gr.Dropdown(
+                    choices=[
+                        "small-obj365", "medium-obj365", "large-obj365",
+                        "small-coco", "medium-coco", "large-coco",
+                        "small-obj2coco", "medium-obj2coco", "large-obj2coco",
+                    ],
+                    value="medium-obj2coco",
+                    label="D-FINE model",
+                )
+
+                objdet_threshold_slider = gr.Slider(
+                    minimum=0.05,
+                    maximum=0.9,
+                    value=0.3,
+                    step=0.05,
+                    label="Detection threshold",
+                )
+
+                btn_objdet = gr.Button(
+                    "Run Object Detection",
+                    variant="primary",
+                )
+
+            with gr.Column(scale=1):
+
+                out_objdet_image = gr.Image(
+                    label="Detections (one colour-coded box per object, class + score)",
+                    height=IMG_HEIGHT,
+                )
+
+                out_objdet_status = gr.Textbox(
+                    label="Detected classes",
+                    lines=12,
+                    interactive=False,
+                )
+
+        btn_objdet.click(
+            fn=run_objdet,
+            inputs=[inp_objdet, objdet_model_dropdown, objdet_threshold_slider],
+            outputs=[out_objdet_image, out_objdet_status],
+            concurrency_limit=1,
+        )
+
 
 app.launch(
     server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),

dfine_jina_pipeline.py

@@ -5,6 +5,7 @@ Outputs jina_crops folder and results CSV.
 
 import argparse
 import csv
+import os
 import time
 from pathlib import Path
 
@@ -776,5 +777,112 @@ def run_single_image(
     return group_crop_images, known_crop_composites, log_text
 
 
+# -----------------------------------------------------------------------------
+# Full-frame object detection (Object Detection tab): run a chosen D-FINE model
+# on the whole image, draw every detection, and summarise the classes found.
+# -----------------------------------------------------------------------------
+
+# Cycled by class id so different classes get distinct box colours.
+OBJDET_PALETTE = [
+    (255, 56, 56), (255, 159, 0), (255, 221, 0), (0, 199, 89),
+    (0, 162, 255), (170, 0, 255), (0, 255, 200), (255, 0, 170),
+    (120, 200, 0), (0, 120, 255),
+]
+
+
+def _objdet_font(size):
+    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
+              "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
+        if os.path.exists(p):
+            from PIL import ImageFont
+            return ImageFont.truetype(p, size)
+    from PIL import ImageFont
+    return ImageFont.load_default()
+
+
+def run_object_detection(
+    pil_image,
+    dfine_model="medium-obj2coco",
+    det_threshold=0.3,
+    device=None,
+):
+    """Run a chosen D-FINE model on the full frame and draw every detection.
+
+    Used by the Object Detection tab. Returns ``(annotated_image_np, status_text)``:
+      - the image with one colour-coded box per detection (label + score),
+      - a text summary of how many of each class were found.
+    """
+    from PIL import Image, ImageDraw
+    import numpy as np
+
+    if pil_image is None:
+        return None, "Upload an image."
+
+    dfine_model = (dfine_model or "medium-obj2coco").strip().lower()
+    if dfine_model not in DFINE_MODEL_IDS:
+        dfine_model = "medium-obj2coco"
+    model_id = DFINE_MODEL_IDS[dfine_model]
+
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    pil = pil_image.convert("RGB") if isinstance(pil_image, Image.Image) else Image.fromarray(pil_image).convert("RGB")
+
+    # Load / reuse D-FINE (shares the app-wide cache with the Detect & Classify tab).
+    global _APP_DFINE
+    if _APP_DFINE is None or _APP_DFINE[0] != dfine_model:
+        print(f"[*] Loading D-FINE ({model_id})...")
+        image_processor = AutoImageProcessor.from_pretrained(model_id)
+        dfine_model_obj = DFineForObjectDetection.from_pretrained(model_id)
+        dfine_model_obj = dfine_model_obj.to(device).eval()
+        person_car_ids = get_person_car_label_ids(dfine_model_obj)
+        _APP_DFINE = (dfine_model, image_processor, dfine_model_obj, person_car_ids)
+
+    _model_id, image_processor, dfine_model_obj, _person_car_ids = _APP_DFINE
+
+    # Use the model's actual device (it may have been loaded on GPU by another tab).
+    mdev = str(next(dfine_model_obj.parameters()).device)
+    detections = run_dfine(pil, image_processor, dfine_model_obj, mdev, det_threshold)
+    detections.sort(key=lambda d: -d["conf"])
+
+    out = pil.copy()
+    draw = ImageDraw.Draw(out)
+    W, H = out.size
+    font = _objdet_font(max(14, int(0.018 * max(W, H))))
+    lw = max(2, int(0.003 * max(W, H)))
+
+    for d in detections:
+        x1, y1, x2, y2 = [int(v) for v in d["box"]]
+        color = OBJDET_PALETTE[d["cls"] % len(OBJDET_PALETTE)]
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=lw)
+        text = f"{d['label']} {d['conf']:.2f}"
+        tb = draw.textbbox((0, 0), text, font=font)
+        tw, th = tb[2] - tb[0], tb[3] - tb[1]
+        ty = max(0, y1 - th - 4)
+        draw.rectangle([x1, ty, x1 + tw + 6, ty + th + 4], fill=color)
+        draw.text((x1 + 3, ty + 2), text, fill=(255, 255, 255), font=font)
+
+    if not detections:
+        return np.array(out), (
+            f"No objects detected at threshold {det_threshold:.2f}. "
+            "Try lowering the threshold."
+        )
+
+    # Summary: which classes were detected in the image, and how many of each.
+    counts = {}
+    for d in detections:
+        counts[d["label"]] = counts.get(d["label"], 0) + 1
+
+    lines = [
+        f"{len(detections)} detection(s) at threshold {det_threshold:.2f} "
+        f"across {len(counts)} class(es):",
+        "",
+    ]
+    lines += [f"  {name}: {n}" for name, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))]
+    lines += ["", "Per detection (sorted by score):"]
+    lines += [f"  {d['label']} {d['conf']:.3f}" for d in detections]
+
+    return np.array(out), "\n".join(lines)
+
+
 if __name__ == "__main__":
     main()