app.py

""" 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, run_object_detection
from ppe_compliance import run_ppe_compliance

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

DEFAULT_LABELS = "gun, knife, cigarette, phone"


def run_dfine_classify(image, dfine_threshold, dfine_model_choice, classifier_choice, siglip_threshold, labels_text):
    """D-FINE first, then classify crops with SigLIP.
    Returns (group_crop_gallery, known_crop_gallery, status_message).
    """
    if image is None:
        return [], [], "Upload an image."

    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
    if not labels:
        return [], [], "Enter at least one label."

    dfine_model = dfine_model_choice.strip().lower() if dfine_model_choice else "medium-obj2coco"
    conf_thresh = float(siglip_threshold)
    classifier = classifier_choice.strip() if classifier_choice else "siglip-256"

    group_crops, known_crops, status = run_single_image(
        image,
        dfine_model=dfine_model,
        det_threshold=float(dfine_threshold),
        conf_threshold=conf_thresh,
        gap_threshold=0.0,
        min_side=24,
        crop_dedup_iou=0.4,
        min_display_conf=conf_thresh,
        classifier=classifier,
        labels=labels,
    )

    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


with gr.Blocks(title="Small Object Detection") as app:

    gr.Markdown("# Small Object Detection")

    with gr.Tabs():

      with gr.Tab("Detect & Classify"):

        gr.Markdown(
            "**D-FINE** detects persons/cars, then small-object crops are classified with **SigLIP** (zero-shot). "
            "Choose a D-FINE model and enter comma-separated class labels for SigLIP."
        )

        with gr.Row():

            with gr.Column(scale=1):

                inp_dfine = gr.Image(
                    type="pil",
                    label="Input image",
                    height=IMG_HEIGHT
                )

                dfine_model_radio = 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",
                )

                classifier_dropdown = gr.Dropdown(
                    choices=["siglip-224", "siglip-256", "siglip-384"],
                    value="siglip-256",
                    label="Classifier model",
                )

                dfine_threshold_slider = gr.Slider(
                    minimum=0.05,
                    maximum=0.5,
                    value=0.15,
                    step=0.05,
                    label="D-FINE detection threshold",
                )

                def update_dfine_threshold_default(choice):
                    if not choice:
                        return gr.update(value=0.15)
                    size = choice.strip().lower().split("-")[0]
                    defaults = {"large": 0.2, "medium": 0.15, "small": 0.1}
                    return gr.update(value=defaults.get(size, 0.15))

                dfine_model_radio.change(
                    fn=update_dfine_threshold_default,
                    inputs=[dfine_model_radio],
                    outputs=[dfine_threshold_slider],
                )

                siglip_threshold_slider = gr.Slider(
                    minimum=0.001,
                    maximum=0.1,
                    value=0.005,
                    step=0.001,
                    label="SigLIP: min confidence threshold",
                )

                labels_input = gr.Textbox(
                    label="Labels (comma-separated)",
                    value=DEFAULT_LABELS,
                    placeholder="e.g. gun, knife, cigarette, phone",
                )

                btn_dfine = gr.Button(
                    "Run D-FINE + Classify",
                    variant="primary"
                )

            with gr.Column(scale=1):

                out_gallery_dfine = gr.Gallery(
                    label="Person/car crops (all D-FINE objects inside drawn with label + score)",
                    height=IMG_HEIGHT,
                    columns=2,
                    object_fit="contain",
                )

                out_gallery_known = gr.Gallery(
                    label="Known objects (class + score above each crop)",
                    height=IMG_HEIGHT,
                    columns=4,
                    object_fit="contain",
                )

                out_status_dfine = gr.Textbox(
                    label="Classification details",
                    lines=8,
                    interactive=False,
                )

        btn_dfine.click(
            fn=run_dfine_classify,
            inputs=[inp_dfine, dfine_threshold_slider, dfine_model_radio, classifier_dropdown, siglip_threshold_slider, labels_input],
            outputs=[out_gallery_dfine, out_gallery_known, out_status_dfine],
            concurrency_limit=1,
        )

      with gr.Tab("PPE Compliance"):

        gr.Markdown(
            "A **D-FINE person detector** (same model as the Detect & Classify tab) finds each "
            "person, then the **fine-tuned D-FINE-M PPE detector** runs on each person crop. "
            "Any of the 6 required items (goggles, helmet, mask, shoes, vest, glove) that are "
            "missing for a person are flagged. **Green** box = compliant, **red** box = violation; "
            "a per-person checklist (1 = present, 0 = missing) is drawn beside each box."
        )

        with gr.Row():

            with gr.Column(scale=1):

                inp_ppe = gr.Image(
                    type="pil",
                    label="Input image",
                    height=IMG_HEIGHT,
                )

                person_threshold_slider = gr.Slider(
                    minimum=0.05,
                    maximum=0.95,
                    value=0.75,
                    step=0.05,
                    label="Person detection threshold",
                )

                ppe_threshold_slider = gr.Slider(
                    minimum=0.05,
                    maximum=0.95,
                    value=0.25,
                    step=0.05,
                    label="PPE detection threshold",
                )

                btn_ppe = gr.Button(
                    "Run PPE Compliance",
                    variant="primary",
                )

            with gr.Column(scale=1):

                out_ppe_image = gr.Image(
                    label="Overview (green = OK, red = violation; #N per person)",
                    height=IMG_HEIGHT,
                )

                out_ppe_status = gr.Textbox(
                    label="Per-person report",
                    lines=10,
                    interactive=False,
                )

        out_ppe_gallery = gr.Gallery(
            label="Per-person breakdown (click a person — PPE worn + checklist: 1 = present, 0 = missing)",
            columns=3,
            height=IMG_HEIGHT,
            object_fit="contain",
        )

        btn_ppe.click(
            fn=lambda image, p_thr, ppe_thr: run_ppe_compliance(
                image,
                person_threshold=float(p_thr),
                ppe_threshold=float(ppe_thr),
            ),
            inputs=[inp_ppe, person_threshold_slider, ppe_threshold_slider],
            outputs=[out_ppe_image, out_ppe_gallery, out_ppe_status],
            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"),
    server_port=int(
        os.environ.get(
            "PORT",
            os.environ.get("GRADIO_SERVER_PORT", 7860)
        )
    ),
)