Add YOLO object detection workflow

app.py

@@ -1,4 +1,4 @@
-"""Hugging Face Space for semantic segmentation of street-scene images."""
+"""Hugging Face Space for street-scene detection and segmentation."""
 
 from __future__ import annotations
 
@@ -34,9 +34,17 @@ from segmentation_utils import (
     resize_for_output,
     write_class_csv,
 )
+from detection_utils import (
+    build_detection_summary,
+    build_detection_table,
+    build_street_indicators,
+    render_detection,
+    write_detection_csv,
+)
 
-MODEL_ID = "nvidia/segformer-b0-finetuned-cityscapes-1024-1024"
-OUTPUT_ROOT = Path("/tmp/street-scene-segmentation")
+SEGMENTATION_MODEL_ID = "nvidia/segformer-b0-finetuned-cityscapes-1024-1024"
+DETECTION_MODEL_ID = "yolo26s.pt"
+OUTPUT_ROOT = Path("/tmp/street-scene-vision")
 SAMPLE_ROOT = (
     "https://raw.githubusercontent.com/"
     "LabMingzeChen/HNIVision/main/space/examples"
@@ -49,9 +57,9 @@ def load_model():
     """Download once per container, then reuse the processor and model."""
     torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+    processor = AutoImageProcessor.from_pretrained(SEGMENTATION_MODEL_ID)
     model = (
-        SegformerForSemanticSegmentation.from_pretrained(MODEL_ID)
+        SegformerForSemanticSegmentation.from_pretrained(SEGMENTATION_MODEL_ID)
         .to(device)
         .eval()
     )
@@ -59,6 +67,121 @@ def load_model():
     return processor, model, id2label, device
 
 
+@lru_cache(maxsize=1)
+def load_detector():
+    """Download YOLO26-s once per container and reuse it."""
+    from ultralytics import YOLO
+
+    return YOLO(DETECTION_MODEL_ID)
+
+
+def _format_detection_indicators(
+    detections: list[dict[str, object]],
+) -> str:
+    indicators = build_street_indicators(detections)
+    average_confidence = (
+        sum(float(item["confidence"]) for item in detections) / len(detections)
+        if detections
+        else 0.0
+    )
+    return f"""
+### Detection-based street indicators
+
+| Indicator | Visible count |
+|---|---:|
+| People | {indicators['people']} |
+| Active-mobility objects (`person` + `bicycle`) | {indicators['active_mobility']} |
+| Motor vehicles | {indicators['motor_vehicles']} |
+| All transport objects | {indicators['all_transport']} |
+| All detected objects | {len(detections)} |
+
+Average detection confidence: **{average_confidence:.2f}**
+
+> Counts describe visible COCO detections in this image. They are not traffic-flow,
+> occupancy, accessibility, or safety measurements.
+"""
+
+
+@spaces.GPU(duration=90)
+def detect_street_objects(
+    image: Image.Image | None,
+    confidence_threshold: float,
+):
+    """Run YOLO object detection and return visual and tabular outputs."""
+    if image is None:
+        raise gr.Error("Please upload a street-scene image first.")
+
+    started_at = time.perf_counter()
+    prepared_image = resize_for_output(image)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    try:
+        detector = load_detector()
+        predictions = detector.predict(
+            source=np.asarray(prepared_image),
+            conf=float(confidence_threshold),
+            imgsz=1024,
+            device=device,
+            max_det=100,
+            verbose=False,
+        )
+        prediction = predictions[0]
+        names = prediction.names
+        detections: list[dict[str, object]] = []
+        if prediction.boxes is not None:
+            coordinates = prediction.boxes.xyxy.detach().cpu().tolist()
+            confidences = prediction.boxes.conf.detach().cpu().tolist()
+            class_ids = prediction.boxes.cls.detach().cpu().tolist()
+            for coordinates_row, confidence, class_id_value in zip(
+                coordinates,
+                confidences,
+                class_ids,
+            ):
+                class_id = int(class_id_value)
+                detections.append(
+                    {
+                        "class_id": class_id,
+                        "class_name": str(names[class_id]),
+                        "confidence": float(confidence),
+                        "x1": float(coordinates_row[0]),
+                        "y1": float(coordinates_row[1]),
+                        "x2": float(coordinates_row[2]),
+                        "y2": float(coordinates_row[3]),
+                    }
+                )
+    except Exception as exc:
+        raise gr.Error(
+            f"Object detection failed: {type(exc).__name__}: {exc}"
+        ) from exc
+
+    overlay = render_detection(prepared_image, detections)
+    summary_rows = build_detection_summary(detections)
+    detection_rows = build_detection_table(detections)
+
+    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
+    output_dir.mkdir(parents=True, exist_ok=True)
+    overlay_path = output_dir / "street_object_detection_overlay.png"
+    csv_path = output_dir / "street_object_detections.csv"
+    overlay.save(overlay_path)
+    write_detection_csv(csv_path, detection_rows)
+
+    elapsed = time.perf_counter() - started_at
+    visible_classes = len(summary_rows)
+    status = (
+        f"Done · {prepared_image.width}×{prepared_image.height} · "
+        f"{len(detections)} objects · {visible_classes} COCO classes · "
+        f"{elapsed:.1f}s · device={device}"
+    )
+    return (
+        overlay,
+        summary_rows,
+        detection_rows,
+        [str(overlay_path), str(csv_path)],
+        _format_detection_indicators(detections),
+        status,
+    )
+
+
 @spaces.GPU(duration=90)
 def segment_street_scene(
     image: Image.Image | None,
@@ -148,16 +271,17 @@ CSS = """
 """
 
 
-with gr.Blocks(title="Street Scene Segmentation", theme=gr.themes.Soft(), css=CSS) as demo:
+with gr.Blocks(title="Street Scene Vision Toolkit", theme=gr.themes.Soft(), css=CSS) as demo:
     gr.Markdown(
         """
         <div class="hero">
-          <h1>🚦 Street Scene Semantic Segmentation</h1>
-          <p>Turn a street image into a pixel-level map of roads, sidewalks, buildings, vegetation, vehicles, people, and more.</p>
-          <p class="muted">SegFormer-B0 · 19 Cityscapes classes · browser-based · no API key required</p>
+          <h1>🚦 Street Scene Vision Toolkit</h1>
+          <p>Detect individual objects with bounding boxes and map every pixel with semantic segmentation.</p>
+          <p class="muted">YOLO26-s · COCO 80 objects · SegFormer-B0 · 19 Cityscapes classes · no API key required</p>
           <div class="project-links">
             <a class="project-link" href="https://huggingface.co/spaces/Mingze/StreetSceneSegmentation" target="_blank">🤗 Hugging Face Space</a>
-            <a class="project-link" href="https://huggingface.co/nvidia/segformer-b0-finetuned-cityscapes-1024-1024" target="_blank">🧠 Model card</a>
+            <a class="project-link" href="https://docs.ultralytics.com/models/yolo26/" target="_blank">📦 YOLO26</a>
+            <a class="project-link" href="https://huggingface.co/nvidia/segformer-b0-finetuned-cityscapes-1024-1024" target="_blank">🎨 SegFormer</a>
             <a class="project-link" href="https://github.com/LabMingzeChen/StreetSceneSegmentation" target="_blank">⭐ GitHub source</a>
           </div>
         </div>
@@ -179,7 +303,15 @@ with gr.Blocks(title="Street Scene Segmentation", theme=gr.themes.Soft(), css=CS
                     label="Try the UBC campus street example",
                     examples_per_page=1,
                 )
-            with gr.Accordion("Display settings", open=False):
+            with gr.Accordion("Object-detection settings", open=True):
+                confidence_input = gr.Slider(
+                    0.05,
+                    0.90,
+                    value=0.25,
+                    step=0.05,
+                    label="Minimum detection confidence",
+                )
+            with gr.Accordion("Segmentation settings", open=False):
                 opacity_input = gr.Slider(
                     0.15,
                     0.85,
@@ -195,47 +327,79 @@ with gr.Blocks(title="Street Scene Segmentation", theme=gr.themes.Soft(), css=CS
                     label="Minimum class area shown in table (%)",
                 )
             with gr.Row():
-                run_button = gr.Button("Run segmentation", variant="primary", size="lg")
-                clear_button = gr.ClearButton(value="Clear", components=[image_input])
-            status_output = gr.Markdown()
+                detection_button = gr.Button("Detect objects", variant="primary", size="lg")
+                segmentation_button = gr.Button("Segment pixels", size="lg")
+            clear_button = gr.ClearButton(value="Clear image", components=[image_input])
 
         with gr.Column(scale=7):
             with gr.Tabs():
-                with gr.Tab("Overlay"):
+                with gr.Tab("Object detection"):
+                    detection_output = gr.Image(
+                        label="YOLO26-s bounding boxes",
+                        height=470,
+                    )
+                    detection_status = gr.Markdown()
+                with gr.Tab("Segmentation overlay"):
                     overlay_output = gr.Image(label="Segmentation overlay", height=470)
+                    segmentation_status = gr.Markdown()
                 with gr.Tab("Color mask"):
                     mask_output = gr.Image(label="Cityscapes color mask", height=470)
 
-    table_output = gr.Dataframe(
-        headers=["Class ID", "Class", "Pixels", "Area share (%)", "Color"],
-        datatype=["number", "str", "number", "number", "str"],
-        label="Detected street-scene classes",
-        interactive=False,
-        wrap=True,
-    )
-    files_output = gr.File(
-        label="Download results (overlay, color mask, class IDs, and CSV)",
-        file_count="multiple",
-    )
+    with gr.Tabs():
+        with gr.Tab("Detection results"):
+            detection_indicators = gr.Markdown()
+            detection_summary = gr.Dataframe(
+                headers=["Class", "Count", "Average confidence", "Maximum confidence"],
+                datatype=["str", "number", "number", "number"],
+                label="Detected object classes",
+                interactive=False,
+                wrap=True,
+            )
+            with gr.Accordion("Detailed bounding-box coordinates", open=False):
+                detection_table = gr.Dataframe(
+                    headers=["Object ID", "Class", "Confidence", "x1", "y1", "x2", "y2"],
+                    datatype=["number", "str", "number", "number", "number", "number", "number"],
+                    label="Individual detections",
+                    interactive=False,
+                    wrap=True,
+                )
+            detection_files = gr.File(
+                label="Download detection overlay and bounding-box CSV",
+                file_count="multiple",
+            )
+
+        with gr.Tab("Segmentation results"):
+            table_output = gr.Dataframe(
+                headers=["Class ID", "Class", "Pixels", "Area share (%)", "Color"],
+                datatype=["number", "str", "number", "number", "str"],
+                label="Detected street-scene classes",
+                interactive=False,
+                wrap=True,
+            )
+            files_output = gr.File(
+                label="Download segmentation overlay, color mask, class IDs, and CSV",
+                file_count="multiple",
+            )
 
     gr.Markdown(
         """
         ## How to use the app
 
         <div class="guide-grid">
-          <div class="guide-card"><h3>1 · Choose an image</h3><p>Upload, paste, use a webcam, or select one of the sample street scenes above.</p></div>
-          <div class="guide-card"><h3>2 · Run segmentation</h3><p>The model assigns one Cityscapes class to every pixel. The first run may take longer while the model loads.</p></div>
-          <div class="guide-card"><h3>3 · Explore and download</h3><p>Compare the overlay and mask, review area shares, and download reusable PNG and CSV outputs.</p></div>
+          <div class="guide-card"><h3>1 · Choose an image</h3><p>Upload, paste, use a webcam, or select the UBC campus example.</p></div>
+          <div class="guide-card"><h3>2 · Choose a method</h3><p>Run YOLO object detection, SegFormer semantic segmentation, or run both methods on the same image.</p></div>
+          <div class="guide-card"><h3>3 · Explore and download</h3><p>Compare boxes, overlays, masks, counts, pixel shares, coordinates, and reusable CSV outputs.</p></div>
         </div>
 
         ## What the results mean
 
-        - **Overlay** blends the prediction with the original photograph. White lines mark boundaries between predicted classes.
-        - **Color mask** shows only the standard Cityscapes colors, making classes easier to compare across images.
-        - **Area share** is the percentage of image pixels assigned to each class. It describes visual coverage, not physical land area.
-        - **Class-ID PNG** stores the numeric class at each pixel; the CSV summarizes visible classes and their colors.
+        - **Object detection** finds separate COCO objects, draws bounding boxes, and reports a confidence score for each detection.
+        - **Detection indicators** summarize visible people, active-mobility objects, and transport objects. They are transparent image counts, not traffic-flow estimates.
+        - **Segmentation overlay** blends the Cityscapes prediction with the original photograph. White lines mark class boundaries.
+        - **Color mask and area share** show pixel-level scene composition. Area share describes visual coverage, not physical land area.
+        - **Downloadable data** include bounding-box coordinates, class-ID pixels, overlays, masks, and CSV summaries.
 
-        | Scene layer | Classes |
+        | Scene layer | Segmentation classes |
         |---|---|
         | Travel surfaces | road, sidewalk |
         | Built environment | building, wall, fence, pole, traffic light, traffic sign |
@@ -245,20 +409,37 @@ with gr.Blocks(title="Street Scene Segmentation", theme=gr.themes.Soft(), css=CS
 
         ## Classroom and research ideas
 
-        - Compare vegetation and built-environment shares across several streets.
-        - Discuss why a pixel classifier can confuse terrain, sidewalk, wall, or fence at object boundaries.
-        - Export the CSV files and create a class-coverage chart for an urban-design exercise.
+        - Compare detected people, bicycles, and motor vehicles across several street images.
+        - Compare what bounding boxes reveal with what pixel-level segmentation reveals.
+        - Discuss missed objects, false positives, confidence thresholds, and segmentation boundary errors.
+        - Export both CSV files and build object-count and class-coverage charts.
         - Compare the same location across seasons, weather conditions, or camera viewpoints.
 
-        > **Important:** predictions are model estimates, not ground truth. The model was trained for road-driving imagery and can be less reliable for unusual viewpoints, severe weather, night scenes, indoor images, or places unlike the training data. Do not use it for safety-critical decisions or to identify individuals.
+        > **Important:** predictions are model estimates, not ground truth. COCO detection is limited to its trained object vocabulary, while Cityscapes segmentation is specialized for road-driving imagery. Do not use either output for safety-critical decisions, surveillance, or identifying individuals.
 
+        [Read the YOLO26 documentation](https://docs.ultralytics.com/models/yolo26/) ·
+        [Explore the COCO dataset](https://cocodataset.org/) ·
         [Read the SegFormer paper](https://arxiv.org/abs/2105.15203) ·
         [Explore the Cityscapes dataset](https://www.cityscapes-dataset.com/) ·
         [View the source on GitHub](https://github.com/LabMingzeChen/StreetSceneSegmentation)
         """
     )
 
-    run_button.click(
+    detection_button.click(
+        fn=detect_street_objects,
+        inputs=[image_input, confidence_input],
+        outputs=[
+            detection_output,
+            detection_summary,
+            detection_table,
+            detection_files,
+            detection_indicators,
+            detection_status,
+        ],
+        api_name="detect",
+    )
+
+    segmentation_button.click(
         fn=segment_street_scene,
         inputs=[image_input, opacity_input, min_share_input],
         outputs=[
@@ -266,7 +447,7 @@ with gr.Blocks(title="Street Scene Segmentation", theme=gr.themes.Soft(), css=CS
             mask_output,
             table_output,
             files_output,
-            status_output,
+            segmentation_status,
         ],
         api_name="segment",
     )

detection_utils.py

@@ -0,0 +1,143 @@
+"""Rendering and summary helpers for street-scene object detection."""
+
+from __future__ import annotations
+
+import csv
+from collections import defaultdict
+from pathlib import Path
+from typing import Iterable
+
+from PIL import Image, ImageDraw, ImageFont
+
+
+STREET_OBJECT_GROUPS: dict[str, set[str]] = {
+    "people": {"person"},
+    "active_mobility": {"person", "bicycle"},
+    "motor_vehicles": {"car", "motorcycle", "bus", "truck", "train"},
+    "all_transport": {
+        "bicycle",
+        "car",
+        "motorcycle",
+        "bus",
+        "truck",
+        "train",
+    },
+}
+
+
+def _load_font(size: int) -> ImageFont.ImageFont:
+    candidates = (
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
+        "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
+        "/System/Library/Fonts/Helvetica.ttc",
+    )
+    for candidate in candidates:
+        try:
+            return ImageFont.truetype(candidate, size)
+        except (OSError, IOError):
+            continue
+    return ImageFont.load_default()
+
+
+def detection_color(class_id: int) -> tuple[int, int, int]:
+    """Return a stable, high-contrast color for a COCO class ID."""
+    return (
+        int((67 * class_id + 37) % 190 + 40),
+        int((97 * class_id + 71) % 190 + 40),
+        int((43 * class_id + 113) % 190 + 40),
+    )
+
+
+def render_detection(
+    image: Image.Image,
+    detections: Iterable[dict[str, object]],
+) -> Image.Image:
+    """Draw labeled bounding boxes on a copy of the input image."""
+    rendered = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(rendered)
+    short_side = min(rendered.size)
+    line_width = max(2, round(short_side / 320))
+    font = _load_font(max(13, min(24, round(short_side / 55))))
+
+    for detection in detections:
+        class_id = int(detection["class_id"])
+        color = detection_color(class_id)
+        box = [
+            float(detection["x1"]),
+            float(detection["y1"]),
+            float(detection["x2"]),
+            float(detection["y2"]),
+        ]
+        label = (
+            f"{detection['class_name']} "
+            f"{float(detection['confidence']):.2f}"
+        )
+        draw.rectangle(box, outline=color, width=line_width)
+        text_box = draw.textbbox((box[0], box[1]), label, font=font)
+        text_height = text_box[3] - text_box[1]
+        text_width = text_box[2] - text_box[0]
+        text_y = max(0.0, box[1] - text_height - 8)
+        background = [box[0], text_y, box[0] + text_width + 8, text_y + text_height + 8]
+        draw.rectangle(background, fill=color)
+        draw.text((box[0] + 4, text_y + 4), label, fill="white", font=font)
+
+    return rendered
+
+
+def build_detection_summary(
+    detections: Iterable[dict[str, object]],
+) -> list[list[object]]:
+    """Aggregate detection counts and confidence by class."""
+    grouped: dict[str, list[float]] = defaultdict(list)
+    for detection in detections:
+        grouped[str(detection["class_name"])].append(
+            float(detection["confidence"])
+        )
+
+    rows = [
+        [class_name, len(confidences), round(sum(confidences) / len(confidences), 3), round(max(confidences), 3)]
+        for class_name, confidences in grouped.items()
+    ]
+    rows.sort(key=lambda row: (-int(row[1]), str(row[0])))
+    return rows
+
+
+def build_detection_table(
+    detections: Iterable[dict[str, object]],
+) -> list[list[object]]:
+    """Build one exportable row per bounding box."""
+    rows: list[list[object]] = []
+    for index, detection in enumerate(detections, start=1):
+        rows.append(
+            [
+                index,
+                str(detection["class_name"]),
+                round(float(detection["confidence"]), 3),
+                round(float(detection["x1"]), 1),
+                round(float(detection["y1"]), 1),
+                round(float(detection["x2"]), 1),
+                round(float(detection["y2"]), 1),
+            ]
+        )
+    return rows
+
+
+def build_street_indicators(
+    detections: Iterable[dict[str, object]],
+) -> dict[str, int]:
+    """Derive transparent street-scene counts from visible COCO objects."""
+    class_names = [str(detection["class_name"]) for detection in detections]
+    return {
+        name: sum(class_name in members for class_name in class_names)
+        for name, members in STREET_OBJECT_GROUPS.items()
+    }
+
+
+def write_detection_csv(
+    path: Path,
+    rows: Iterable[Iterable[object]],
+) -> None:
+    with path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.writer(handle)
+        writer.writerow(["object_id", "class", "confidence", "x1", "y1", "x2", "y2"])
+        writer.writerows(rows)

requirements.txt

@@ -3,3 +3,4 @@ transformers==4.57.6
 torch==2.8.0
 numpy==2.2.6
 Pillow==11.3.0
+ultralytics==8.4.102

tests/test_detection_utils.py

@@ -0,0 +1,57 @@
+from PIL import Image
+
+import detection_utils as utils
+
+
+DETECTIONS = [
+    {
+        "class_id": 0,
+        "class_name": "person",
+        "confidence": 0.9,
+        "x1": 1.0,
+        "y1": 2.0,
+        "x2": 10.0,
+        "y2": 20.0,
+    },
+    {
+        "class_id": 2,
+        "class_name": "car",
+        "confidence": 0.7,
+        "x1": 12.0,
+        "y1": 4.0,
+        "x2": 30.0,
+        "y2": 18.0,
+    },
+    {
+        "class_id": 2,
+        "class_name": "car",
+        "confidence": 0.5,
+        "x1": 32.0,
+        "y1": 5.0,
+        "x2": 45.0,
+        "y2": 17.0,
+    },
+]
+
+
+def test_detection_summary_groups_and_sorts_classes():
+    assert utils.build_detection_summary(DETECTIONS) == [
+        ["car", 2, 0.6, 0.7],
+        ["person", 1, 0.9, 0.9],
+    ]
+
+
+def test_street_indicators_are_transparent_counts():
+    assert utils.build_street_indicators(DETECTIONS) == {
+        "people": 1,
+        "active_mobility": 1,
+        "motor_vehicles": 2,
+        "all_transport": 2,
+    }
+
+
+def test_detection_overlay_matches_input_size():
+    image = Image.new("RGB", (60, 40), "white")
+    rendered = utils.render_detection(image, DETECTIONS)
+    assert rendered.size == image.size
+    assert rendered.getpixel((1, 2)) != (255, 255, 255)