Cooked I Am

pipeline/ops/__init__.py

@@ -0,0 +1,41 @@
+"""Operations and utility functions."""
+
+from pipeline.ops.drawing import (
+    draw_alert_banner,
+    draw_box,
+    draw_detections,
+    draw_fps,
+    draw_frame_overlay,
+    draw_label,
+    draw_zone,
+    get_color,
+)
+from pipeline.ops.geometry import (
+    box_center,
+    box_in_zone,
+    expand_box,
+    point_in_box,
+    point_in_polygon,
+)
+from pipeline.ops.nms import nms, nms_by_class
+
+__all__ = [
+    # NMS
+    "nms",
+    "nms_by_class",
+    # Geometry
+    "expand_box",
+    "point_in_box",
+    "box_center",
+    "point_in_polygon",
+    "box_in_zone",
+    # Drawing
+    "get_color",
+    "draw_box",
+    "draw_label",
+    "draw_detections",
+    "draw_frame_overlay",
+    "draw_alert_banner",
+    "draw_zone",
+    "draw_fps",
+]

pipeline/ops/drawing.py

@@ -0,0 +1,117 @@
+"""Drawing utilities for visualization."""
+
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+
+from pipeline.types import Detection
+
+# Color palette (BGR)
+COLORS = {
+    "person": (255, 0, 0),  # Blue
+    "candidate": (0, 255, 255),  # Yellow
+    "verified": (0, 0, 255),  # Red
+    "default": (255, 255, 255),  # White
+}
+
+
+def get_color(label: str) -> Tuple[int, int, int]:
+    """Get color for a detection label."""
+    for key, color in COLORS.items():
+        if key in label:
+            return color
+    return COLORS["default"]
+
+
+def draw_box(
+    image: np.ndarray,
+    box: Tuple[int, int, int, int],
+    color: Tuple[int, int, int],
+    thickness: int = 2,
+):
+    """Draw a rectangle on image."""
+    x1, y1, x2, y2 = box
+    cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
+
+
+def draw_label(
+    image: np.ndarray,
+    text: str,
+    position: Tuple[int, int],
+    color: Tuple[int, int, int],
+    font_scale: float = 0.5,
+):
+    """Draw a label with background."""
+    x, y = position
+    (tw, _), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1)
+    cv2.rectangle(image, (x, y - 20), (x + tw, y), color, -1)
+    cv2.putText(
+        image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), 1
+    )
+
+
+def draw_detections(image: np.ndarray, detections: List[Detection]):
+    """Draw all detections with boxes and labels."""
+    for det in detections:
+        color = get_color(det.label)
+        draw_box(image, (det.box.x1, det.box.y1, det.box.x2, det.box.y2), color)
+        draw_label(
+            image, f"{det.label} {det.score:.2f}", (det.box.x1, det.box.y1), color
+        )
+
+
+def draw_frame_overlay(image: np.ndarray, frame_id: int):
+    """Draw frame counter overlay."""
+    cv2.putText(
+        image,
+        f"Frame: {frame_id}",
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1.0,
+        (0, 255, 0),
+        2,
+    )
+
+
+def draw_alert_banner(image: np.ndarray, text: str = "SMOKING DETECTED"):
+    """Draw a red alert banner at the top."""
+    cv2.rectangle(image, (0, 0), (image.shape[1], 40), (0, 0, 255), -1)
+    cv2.putText(
+        image,
+        text,
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1,
+        (255, 255, 255),
+        2,
+    )
+
+
+def draw_zone(
+    image: np.ndarray,
+    zone: List[Tuple[int, int]],
+    color: Tuple[int, int, int] = (0, 255, 0),
+    alpha: float = 0.3,
+):
+    """Draw a semi-transparent polygon zone."""
+    overlay = image.copy()
+    pts = np.array(zone, np.int32).reshape((-1, 1, 2))
+    cv2.fillPoly(overlay, [pts], color)
+    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
+    cv2.polylines(image, [pts], True, color, 2)
+
+
+def draw_fps(image: np.ndarray, fps: float):
+    """Draw FPS counter in top-right corner."""
+    h, w = image.shape[:2]
+    text = f"FPS: {fps:.1f}"
+    cv2.putText(
+        image,
+        text,
+        (w - 120, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        0.7,
+        (0, 255, 0),
+        2,
+    )

pipeline/ops/geometry.py

@@ -0,0 +1,73 @@
+"""Geometry utilities for bounding box operations."""
+
+from typing import List, Tuple
+
+from pipeline.types import BoundingBox
+
+
+def expand_box(
+    box: BoundingBox,
+    margin: float,
+    image_size: Tuple[int, int],
+) -> Tuple[int, int, int, int]:
+    """
+    Expand a bounding box by a percentage margin, clamped to image bounds.
+
+    Args:
+        box: Original bounding box
+        margin: Expansion factor (0.2 = 20% on each side)
+        image_size: (width, height) of the image
+
+    Returns:
+        (x1, y1, x2, y2) of expanded box
+    """
+    w_img, h_img = image_size
+    bw = box.x2 - box.x1
+    bh = box.y2 - box.y1
+
+    x1 = max(0, int(box.x1 - bw * margin))
+    y1 = max(0, int(box.y1 - bh * margin))
+    x2 = min(w_img, int(box.x2 + bw * margin))
+    y2 = min(h_img, int(box.y2 + bh * margin))
+
+    return x1, y1, x2, y2
+
+
+def point_in_box(point: Tuple[int, int], box: BoundingBox) -> bool:
+    """Check if a point (x, y) is inside a bounding box."""
+    x, y = point
+    return box.x1 <= x <= box.x2 and box.y1 <= y <= box.y2
+
+
+def box_center(box: BoundingBox) -> Tuple[int, int]:
+    """Get center point of a bounding box."""
+    cx = (box.x1 + box.x2) // 2
+    cy = (box.y1 + box.y2) // 2
+    return cx, cy
+
+
+def point_in_polygon(point: Tuple[int, int], polygon: List[Tuple[int, int]]) -> bool:
+    """
+    Ray casting algorithm to check if point is inside polygon.
+    Useful for zone filtering (ROI).
+    """
+    x, y = point
+    n = len(polygon)
+    inside = False
+
+    j = n - 1
+    for i in range(n):
+        xi, yi = polygon[i]
+        xj, yj = polygon[j]
+
+        if ((yi > y) != (yj > y)) and (x < (xj - xi) * (y - yi) / (yj - yi) + xi):
+            inside = not inside
+        j = i
+
+    return inside
+
+
+def box_in_zone(box: BoundingBox, zone: List[Tuple[int, int]]) -> bool:
+    """Check if box center is inside a zone polygon."""
+    center = box_center(box)
+    return point_in_polygon(center, zone)

pipeline/ops/nms.py

@@ -0,0 +1,61 @@
+"""Non-Maximum Suppression algorithms for detection filtering."""
+
+from typing import List
+
+from pipeline.types import Detection
+
+
+def nms(detections: List[Detection], iou_threshold: float = 0.5) -> List[Detection]:
+    """
+    Non-Maximum Suppression: Remove overlapping boxes, keeping the highest confidence ones.
+
+    Args:
+        detections: List of Detection objects
+        iou_threshold: Boxes with IoU > threshold are suppressed
+
+    Returns:
+        Filtered list of detections
+    """
+    if not detections:
+        return []
+
+    # Sort by confidence (highest first)
+    sorted_dets = sorted(detections, key=lambda d: d.score, reverse=True)
+
+    keep = []
+    while sorted_dets:
+        # Take the best one
+        best = sorted_dets.pop(0)
+        keep.append(best)
+
+        # Filter out boxes that overlap too much with 'best'
+        remaining = []
+        for det in sorted_dets:
+            if best.box.iou(det.box) < iou_threshold:
+                remaining.append(det)
+        sorted_dets = remaining
+
+    return keep
+
+
+def nms_by_class(
+    detections: List[Detection], iou_threshold: float = 0.5
+) -> List[Detection]:
+    """
+    Run NMS separately for each class label.
+    This prevents boxes of different classes from suppressing each other.
+    """
+    if not detections:
+        return []
+
+    # Group by label
+    by_label: dict[str, List[Detection]] = {}
+    for det in detections:
+        by_label.setdefault(det.label, []).append(det)
+
+    # Run NMS per class
+    result = []
+    for dets in by_label.values():
+        result.extend(nms(dets, iou_threshold))
+
+    return result