app.py

"""
Object Counter — Gradio App (Python script version)

Detects and counts any object(s) a user names in an uploaded image, using
YOLO-World (open-vocabulary object detection). No fixed class list — type
what you want counted at runtime.

Usage:
    pip install -r requirements.txt
    python app.py

Then open the local URL Gradio prints (usually http://127.0.0.1:7860).
"""

from collections import Counter

import gradio as gr
import supervision as sv
import torch
from ultralytics import YOLOWorld


# ---------------------------------------------------------------------------
# Workaround for a known, still-open Gradio bug (gradio-app/gradio#11722):
# Gradio's automatic API-schema generation crashes with
# "TypeError: argument of type 'bool' is not iterable" when a component's
# JSON schema contains a boolean where a dict is expected. This has hit many
# projects across several Gradio versions and isn't reliably fixed by
# pinning one version, so we patch the offending function defensively.
# This only affects auto-generated API docs, not the app's own behavior.
# ---------------------------------------------------------------------------
def _patch_gradio_bool_schema_bug():
    try:
        import gradio_client.utils as client_utils

        _original_get_type = client_utils.get_type

        def _patched_get_type(schema):
            if isinstance(schema, bool):
                return "bool"
            return _original_get_type(schema)

        client_utils.get_type = _patched_get_type

        if hasattr(client_utils, "_json_schema_to_python_type"):
            _original_json_type = client_utils._json_schema_to_python_type

            def _patched_json_type(schema, defs=None):
                if isinstance(schema, bool):
                    return "bool"
                return _original_json_type(schema, defs)

            client_utils._json_schema_to_python_type = _patched_json_type
    except Exception:
        # If gradio_client's internals change shape in a future version,
        # fail quietly rather than blocking the app from starting.
        pass


_patch_gradio_bool_schema_bug()

# ---------------------------------------------------------------------------
# Model setup
# ---------------------------------------------------------------------------
# "yolov8s-worldv2.pt" is a good balance of speed/accuracy. Swap for
# "yolov8m-worldv2.pt" or "yolov8l-worldv2.pt" if you want more accuracy
# and have the GPU/CPU headroom for it. Weights download automatically
# on first run.
MODEL_WEIGHTS = "yolov8s-worldv2.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

model = YOLOWorld(MODEL_WEIGHTS)
model.to(DEVICE)

box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_scale=0.5, text_thickness=1)

EXAMPLE_PROMPTS = [
    "person, car",
    "car, truck, bus, bicycle",
    "steel rod, stick",
    "bottle, cup, chair",
]


# ---------------------------------------------------------------------------
# Core detection + counting logic
# ---------------------------------------------------------------------------
def count_objects(image, class_text, confidence):
    """
    Run open-vocabulary detection on `image` for the classes in `class_text`
    (comma-separated). Returns:
      - the annotated image
      - a [ [class, count], ... ] table for the results grid
      - an HTML string with a big total-count readout
      - a status message
    """
    if image is None:
        return None, [], _total_html(None), "⚠️ Please upload an image first."
    if not class_text or not class_text.strip():
        return (
            None,
            [],
            _total_html(None),
            "⚠️ Type at least one object to count, e.g. 'car, person'.",
        )

    # Parse comma-separated class names, strip whitespace, drop empties
    classes = [c.strip() for c in class_text.split(",") if c.strip()]

    # Workaround for a known ultralytics/CLIP device-mismatch bug: set_classes()
    # loads a fresh CLIP text encoder that can end up on a different device than
    # the rest of the model when running on GPU. Moving to CPU for this call,
    # then back to the target device, avoids the crash.
    model.to("cpu")
    model.set_classes(classes)
    model.to(DEVICE)

    # Gradio passes images in as RGB numpy arrays already
    results = model.predict(image, conf=confidence, verbose=False)
    result = results[0]

    if len(result.boxes) == 0:
        table = [[cls, 0] for cls in classes]
        status = (
            "No objects detected. Try lowering the confidence threshold "
            "or rewording the class names."
        )
        return image, table, _total_html(0), status

    detections = sv.Detections.from_ultralytics(result)

    class_ids = result.boxes.cls.cpu().numpy().astype(int)
    class_names = [result.names[i] for i in class_ids]
    counts = Counter(class_names)

    labels = [
        f"{result.names[cid]} {conf:.2f}"
        for cid, conf in zip(detections.class_id, detections.confidence)
    ]

    annotated = box_annotator.annotate(scene=image.copy(), detections=detections)
    annotated = label_annotator.annotate(scene=annotated, detections=detections, labels=labels)

    # Table rows, in the order the user typed the classes
    table = [[cls, counts.get(cls, 0)] for cls in classes]
    total = sum(counts.values())
    status = f"✅ Done — found {total} object{'s' if total != 1 else ''}."

    return annotated, table, _total_html(total), status


def _total_html(total):
    """Small styled readout for the total object count."""
    display_value = "—" if total is None else str(total)
    return (
        "<div style='text-align:center; padding: 6px 0;'>"
        "<div style='font-size:0.85rem; letter-spacing:0.05em; text-transform:uppercase; "
        "color:var(--body-text-color-subdued); margin-bottom:2px;'>Total objects</div>"
        f"<div style='font-size:2.4rem; font-weight:700; line-height:1;'>{display_value}</div>"
        "</div>"
    )


def clear_all():
    return None, "", 0.15, None, [], _total_html(None), ""


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
CUSTOM_CSS = """
#header { text-align: center; margin-bottom: 4px; }
#header h1 { margin-bottom: 2px; }
#header p { color: var(--body-text-color-subdued); margin-top: 0; }
#run_btn { min-height: 46px; font-size: 1.05rem; }
#results_table table { font-size: 0.95rem; }
.gr-accordion { margin-top: 8px; }
"""

THEME = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)


def build_app():
    with gr.Blocks(title="Object Counter", theme=THEME, css=CUSTOM_CSS) as demo:
        with gr.Column(elem_id="header"):
            gr.Markdown(
                "# 🔎 Object Counter\n"
                "Upload a photo, tell it what to look for, and get a count — "
                "for **any object**, not just a fixed list."
            )

        with gr.Row(equal_height=False):
            # -------------------- Left: inputs --------------------
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Image",
                    type="numpy",
                    height=340,
                )
                class_input = gr.Textbox(
                    label="What should I count?",
                    placeholder="e.g. car, person, steel rod",
                    info="Separate multiple objects with commas.",
                )
                gr.Examples(
                    examples=EXAMPLE_PROMPTS,
                    inputs=class_input,
                    label="Quick examples",
                )
                confidence_input = gr.Slider(
                    minimum=0.01,
                    maximum=0.9,
                    value=0.15,
                    step=0.01,
                    label="Confidence threshold",
                    info="Lower catches more objects but risks false positives; higher is stricter.",
                )
                with gr.Row():
                    clear_btn = gr.Button("Clear", scale=1)
                    submit_btn = gr.Button("Count objects", variant="primary", scale=2, elem_id="run_btn")

                with gr.Accordion("💡 Tips for better results", open=False):
                    gr.Markdown(
                        "- Be descriptive: **\"steel rod\"** works better than **\"metal\"**.\n"
                        "- Missing objects? Lower the confidence threshold.\n"
                        "- Too many false positives? Raise the threshold or narrow the wording.\n"
                        "- Objects that are tightly packed or overlapping (e.g. a bundle of rods) "
                        "may be undercounted — this is a limitation of box-based detection."
                    )

            # -------------------- Right: outputs --------------------
            with gr.Column(scale=1):
                image_output = gr.Image(label="Detected objects", height=340)
                total_output = gr.HTML(_total_html(None))
                results_table = gr.Dataframe(
                    headers=["Object", "Count"],
                    datatype=["str", "number"],
                    row_count=(0, "dynamic"),
                    col_count=(2, "fixed"),
                    label="Breakdown",
                    elem_id="results_table",
                )
                status_output = gr.Markdown("")

        submit_btn.click(
            fn=count_objects,
            inputs=[image_input, class_input, confidence_input],
            outputs=[image_output, results_table, total_output, status_output],
        )
        clear_btn.click(
            fn=clear_all,
            inputs=None,
            outputs=[
                image_input,
                class_input,
                confidence_input,
                image_output,
                results_table,
                total_output,
                status_output,
            ],
        )

    return demo


if __name__ == "__main__":
    app = build_app()
    app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)