Cooked I Am
This commit is contained in:
1 parent
7b2813be20
commit
d57a48eec6
20 files changed
+665
-276
No files matched your search
Vendored
+1
-1
@@ -5,5 +5,5 @@
|
||||
"**/*.rpymc": true,
|
||||
"**/cache/": true
|
||||
},
|
||||
"pylint.args": ["--generated-members", "cv2.*"],
|
||||
"pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"],
|
||||
}
|
||||
@@ -8,6 +8,8 @@ from pipeline.core import SmokingDetectionPipeline
|
||||
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
|
||||
from pipeline.sinks import (
|
||||
FileAlertStore,
|
||||
MetricsSink,
|
||||
ScreenshotSink,
|
||||
SmartEventSink,
|
||||
VideoFileStore,
|
||||
WindowDisplaySink,
|
||||
@@ -135,8 +137,10 @@ def main():
|
||||
sinks = []
|
||||
video_sink = None
|
||||
smart_sink = None
|
||||
alert_store = None
|
||||
display_sink = None
|
||||
|
||||
# Metrics sink (always active - provides performance summary)
|
||||
metrics_sink = MetricsSink(log_interval=100)
|
||||
sinks.append(metrics_sink)
|
||||
|
||||
if not args.no_save:
|
||||
# Initialize Output Sink
|
||||
@@ -148,6 +152,10 @@ def main():
|
||||
smart_sink = SmartEventSink("events", fps=source.fps)
|
||||
sinks.append(smart_sink)
|
||||
|
||||
# Screenshot sink (saves evidence JPGs)
|
||||
screenshot_sink = ScreenshotSink("screenshots")
|
||||
sinks.append(screenshot_sink)
|
||||
|
||||
# Initialize Alert Store (always active)
|
||||
alert_store = FileAlertStore("alerts.jsonl")
|
||||
sinks.append(alert_store)
|
||||
|
||||
+2
-1
@@ -5,8 +5,9 @@ from collections import deque
|
||||
from concurrent.futures import Future
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from pipeline.ops import nms_by_class
|
||||
from pipeline.protocols import Model, Sink, StreamSource, Verifier
|
||||
from pipeline.types import BoundingBox, Detection, Frame, Ticket, nms_by_class
|
||||
from pipeline.types import BoundingBox, Detection, Frame, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Model wrappers for the pipeline."""
|
||||
|
||||
from pipeline.models.yolo import (
|
||||
DEVICE,
|
||||
YoloCigaretteDetector,
|
||||
YoloClassifier,
|
||||
YoloPersonDetector,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DEVICE",
|
||||
"YoloPersonDetector",
|
||||
"YoloCigaretteDetector",
|
||||
"YoloClassifier",
|
||||
]
|
||||
@@ -1,3 +1,5 @@
|
||||
"""YOLO model wrappers for detection and classification."""
|
||||
|
||||
import logging
|
||||
from typing import Any, List
|
||||
|
||||
@@ -18,12 +20,13 @@ else:
|
||||
|
||||
|
||||
class YoloPersonDetector(Model):
|
||||
def __init__(self, model_path: str = "yolo11n.pt"): # Using small model for speed
|
||||
"""Detects people using YOLO."""
|
||||
|
||||
def __init__(self, model_path: str = "yolo11n.pt"):
|
||||
logger.info("Loading Person Detector: %s", model_path)
|
||||
self.model = YOLO(model_path) # Load model
|
||||
self.model = YOLO(model_path)
|
||||
|
||||
def predict(self, input_data: Any) -> List[Detection]:
|
||||
# input_data: numpy array (H, W, C)
|
||||
results = self.model(
|
||||
input_data, classes=[0], verbose=False, device=DEVICE
|
||||
) # 0 is person
|
||||
@@ -45,13 +48,14 @@ class YoloPersonDetector(Model):
|
||||
|
||||
|
||||
class YoloCigaretteDetector(Model):
|
||||
def __init__(self, model_path: str):
|
||||
"""Detects cigarettes using custom YOLO model."""
|
||||
|
||||
def __init__(self, model_path: str, confidence_threshold: float = 0.2):
|
||||
logger.info("Loading Cigarette Detector: %s", model_path)
|
||||
self.model = YOLO(model_path)
|
||||
self.conf_threshold = confidence_threshold
|
||||
|
||||
def predict(self, input_data: Any) -> List[Detection]:
|
||||
# This model should be trained on cigarette class
|
||||
# Assuming class 0 is cigarette for this custom model
|
||||
results = self.model(input_data, verbose=False, device=DEVICE)
|
||||
|
||||
detections = []
|
||||
@@ -60,8 +64,7 @@ class YoloCigaretteDetector(Model):
|
||||
for box in boxes:
|
||||
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
||||
conf = float(box.conf[0])
|
||||
# Maybe filter by conf here if needed
|
||||
if conf > 0.2:
|
||||
if conf > self.conf_threshold:
|
||||
detections.append(
|
||||
Detection(
|
||||
box=BoundingBox(x1, y1, x2, y2, conf),
|
||||
@@ -73,18 +76,18 @@ class YoloCigaretteDetector(Model):
|
||||
|
||||
|
||||
class YoloClassifier(Verifier):
|
||||
"""Classifies smoking using YOLO classification model."""
|
||||
|
||||
def __init__(self, model_path: str):
|
||||
logger.info("Loading Smoking Classifier: %s", model_path)
|
||||
logger.info("Loading Smoking Classifier: %s", model_path)
|
||||
self.model = YOLO(model_path)
|
||||
|
||||
def verify(self, input_data: Any) -> float:
|
||||
# input_data: numpy array (H, W, C) - Usually a crop
|
||||
if input_data.size == 0:
|
||||
return 0.0
|
||||
|
||||
results = self.model(input_data, verbose=False, device=DEVICE)
|
||||
|
||||
# Check if we got any results
|
||||
if not results:
|
||||
return 0.0
|
||||
|
||||
@@ -92,16 +95,11 @@ class YoloClassifier(Verifier):
|
||||
if probs is None:
|
||||
return 0.0
|
||||
|
||||
# User specified Class 0 is "smoke"
|
||||
# We return the probability of class 0 specifically.
|
||||
# Ultralytics probs.data is a tensor of scores for each class.
|
||||
# Class 0 is "smoke" - return its probability
|
||||
try:
|
||||
return float(probs.data[0])
|
||||
except (AttributeError, IndexError):
|
||||
# Fallback if structure is unexpected, though data[0] is standard for classification
|
||||
logger.warning(
|
||||
"Could not retrieve Class 0 probability. Defaulting to top1conf if class 0."
|
||||
)
|
||||
logger.warning("Could not retrieve Class 0 probability. Checking top1.")
|
||||
if hasattr(probs, "top1") and probs.top1 == 0:
|
||||
return float(probs.top1conf)
|
||||
return 0.0
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -3,10 +3,14 @@
|
||||
from pipeline.sinks.alerts import FileAlertStore
|
||||
from pipeline.sinks.display import WindowDisplaySink
|
||||
from pipeline.sinks.events import SmartEventSink
|
||||
from pipeline.sinks.metrics import MetricsSink
|
||||
from pipeline.sinks.screenshots import ScreenshotSink
|
||||
from pipeline.sinks.video import VideoFileStore
|
||||
|
||||
__all__ = [
|
||||
"FileAlertStore",
|
||||
"MetricsSink",
|
||||
"ScreenshotSink",
|
||||
"SmartEventSink",
|
||||
"VideoFileStore",
|
||||
"WindowDisplaySink",
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
"""Shared drawing utilities for sinks."""
|
||||
|
||||
from typing import List
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from pipeline.types import Detection
|
||||
|
||||
|
||||
def draw_detections(image: np.ndarray, detections: List[Detection]):
|
||||
"""Draw detection boxes with labels and color coding."""
|
||||
for det in detections:
|
||||
x1, y1, x2, y2 = det.box.x1, det.box.y1, det.box.x2, det.box.y2
|
||||
label = det.label
|
||||
score = det.score
|
||||
|
||||
# Color coding
|
||||
color = (255, 255, 255) # White default
|
||||
if "person" in label:
|
||||
color = (255, 0, 0) # Blue
|
||||
elif "candidate" in label:
|
||||
color = (0, 255, 255) # Yellow
|
||||
elif "verified" in label:
|
||||
color = (0, 0, 255) # Red (Confirmed Smoke)
|
||||
|
||||
# Draw Box
|
||||
cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
|
||||
|
||||
# Draw Label
|
||||
text = f"{label} {score:.2f}"
|
||||
(tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
cv2.rectangle(image, (x1, y1 - 20), (x1 + tw, y1), color, -1)
|
||||
cv2.putText(
|
||||
image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
@@ -4,8 +4,8 @@ import logging
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.ops import draw_detections, draw_frame_overlay
|
||||
from pipeline.protocols import Sink
|
||||
from pipeline.sinks.base import draw_detections, draw_frame_overlay
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Metrics sink - tracks pipeline performance stats."""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Dict
|
||||
|
||||
from pipeline.protocols import Sink
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MetricsSink(Sink):
|
||||
"""
|
||||
Tracks pipeline performance metrics:
|
||||
- FPS (frames per second)
|
||||
- Detection counts by label
|
||||
- Alert count
|
||||
- Processing latency
|
||||
"""
|
||||
|
||||
def __init__(self, log_interval: int = 100):
|
||||
self.log_interval = log_interval
|
||||
|
||||
# Timing
|
||||
self.start_time = time.time()
|
||||
self.frame_count = 0
|
||||
self.last_log_time = self.start_time
|
||||
|
||||
# Detection stats
|
||||
self.detection_counts: Dict[str, int] = defaultdict(int)
|
||||
self.alert_count = 0
|
||||
|
||||
# Latency tracking
|
||||
self.frame_timestamps: list[float] = []
|
||||
|
||||
logger.info("MetricsSink initialized. Logging every %d frames.", log_interval)
|
||||
|
||||
def handle(self, frame: Frame):
|
||||
self.frame_count += 1
|
||||
now = time.time()
|
||||
|
||||
# Track latency (time from capture to processing)
|
||||
latency = now - frame.timestamp
|
||||
self.frame_timestamps.append(latency)
|
||||
# Keep only last 100 for rolling average
|
||||
if len(self.frame_timestamps) > 100:
|
||||
self.frame_timestamps.pop(0)
|
||||
|
||||
# Count detections by label
|
||||
for det in frame.detections:
|
||||
self.detection_counts[det.label] += 1
|
||||
|
||||
# Count alerts
|
||||
if frame.ticket:
|
||||
self.alert_count += 1
|
||||
|
||||
# Periodic logging
|
||||
if self.frame_count % self.log_interval == 0:
|
||||
self._log_stats(now)
|
||||
|
||||
def _log_stats(self, now: float):
|
||||
elapsed = now - self.start_time
|
||||
fps = self.frame_count / elapsed if elapsed > 0 else 0
|
||||
avg_latency = (
|
||||
sum(self.frame_timestamps) / len(self.frame_timestamps)
|
||||
if self.frame_timestamps
|
||||
else 0
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[METRICS] Frames: %d | FPS: %.1f | Latency: %.1fms | Alerts: %d",
|
||||
self.frame_count,
|
||||
fps,
|
||||
avg_latency * 1000,
|
||||
self.alert_count,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Print final summary."""
|
||||
elapsed = time.time() - self.start_time
|
||||
fps = self.frame_count / elapsed if elapsed > 0 else 0
|
||||
avg_latency = (
|
||||
sum(self.frame_timestamps) / len(self.frame_timestamps)
|
||||
if self.frame_timestamps
|
||||
else 0
|
||||
)
|
||||
|
||||
logger.info("=" * 50)
|
||||
logger.info("PIPELINE METRICS SUMMARY")
|
||||
logger.info("=" * 50)
|
||||
logger.info("Total Frames: %d", self.frame_count)
|
||||
logger.info("Runtime: %.1f seconds", elapsed)
|
||||
logger.info("Average FPS: %.1f", fps)
|
||||
logger.info("Avg Latency: %.1f ms", avg_latency * 1000)
|
||||
logger.info("Total Alerts: %d", self.alert_count)
|
||||
logger.info("-" * 50)
|
||||
logger.info("Detections by Label:")
|
||||
for label, count in sorted(self.detection_counts.items()):
|
||||
logger.info(" %-25s %d", label, count)
|
||||
logger.info("=" * 50)
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Screenshot sink - saves evidence images on alerts."""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.ops.drawing import draw_alert_banner, draw_detections
|
||||
from pipeline.protocols import Sink
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ScreenshotSink(Sink):
|
||||
"""
|
||||
Saves a screenshot when a ticket (alert) fires.
|
||||
Useful for evidence/audit trails without full video recording.
|
||||
"""
|
||||
|
||||
def __init__(self, output_dir: str = "screenshots", annotate: bool = True):
|
||||
self.output_dir = Path(output_dir)
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.annotate = annotate
|
||||
self.count = 0
|
||||
|
||||
logger.info("ScreenshotSink initialized. Saving to %s", output_dir)
|
||||
|
||||
def handle(self, frame: Frame):
|
||||
if not frame.ticket:
|
||||
return
|
||||
|
||||
self.count += 1
|
||||
ticket = frame.ticket
|
||||
|
||||
# Build filename with timestamp and confidence
|
||||
filename = f"alert_{ticket.frame_id}_{ticket.confidence:.2f}.jpg"
|
||||
filepath = self.output_dir / filename
|
||||
|
||||
# Prepare image
|
||||
if self.annotate:
|
||||
img = frame.data.copy()
|
||||
draw_detections(img, frame.detections)
|
||||
draw_alert_banner(img, f"SMOKING {ticket.confidence:.0%}")
|
||||
else:
|
||||
img = frame.data
|
||||
|
||||
# Save
|
||||
cv2.imwrite(str(filepath), img)
|
||||
logger.info("Saved screenshot: %s", filepath)
|
||||
|
||||
def close(self):
|
||||
logger.info("ScreenshotSink closed. Total screenshots: %d", self.count)
|
||||
@@ -9,8 +9,8 @@ from typing import Optional
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.ops import draw_detections, draw_frame_overlay
|
||||
from pipeline.protocols import Sink
|
||||
from pipeline.sinks.base import draw_detections, draw_frame_overlay
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1,133 +0,0 @@
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from queue import Empty, Queue
|
||||
from typing import Iterator
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.protocols import StreamSource
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileStreamSource(StreamSource):
|
||||
def __init__(self, path: str):
|
||||
self.path = path
|
||||
self._fps = 30.0
|
||||
|
||||
# Probe for FPS if it's a single file
|
||||
if "*" not in self.path:
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.path)
|
||||
if cap.isOpened():
|
||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
cap.release()
|
||||
except Exception as e:
|
||||
logger.warning("Could not probe FPS from %s: %s", path, e)
|
||||
|
||||
@property
|
||||
def fps(self) -> float:
|
||||
return self._fps
|
||||
|
||||
def stream(self) -> Iterator[Frame]:
|
||||
# Handle glob patterns if path is a directory
|
||||
if "*" in self.path:
|
||||
import glob
|
||||
|
||||
files = glob.glob(self.path)
|
||||
# Todo: handle image sequence
|
||||
return
|
||||
|
||||
logger.info("Opening video file: %s", self.path)
|
||||
cap = cv2.VideoCapture(self.path)
|
||||
frame_id = 0
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
logger.info("End of video file reached.")
|
||||
break
|
||||
|
||||
yield Frame(timestamp=time.time(), data=frame, frame_id=frame_id)
|
||||
frame_id += 1
|
||||
cap.release()
|
||||
|
||||
def stop(self):
|
||||
"""No-op for file source (stream ends naturally)."""
|
||||
pass
|
||||
|
||||
|
||||
class ThreadedRTSPSource(StreamSource):
|
||||
"""
|
||||
Reads from RTSP stream or webcam in a background thread to prevent buffer buildup
|
||||
and allow main processing loop to run at its own pace without stalling the decoding.
|
||||
"""
|
||||
|
||||
def __init__(self, source: str | int, queue_size: int = 5):
|
||||
self.source = source
|
||||
self.queue: Queue[Frame] = Queue(maxsize=queue_size)
|
||||
self._fps = 30.0
|
||||
|
||||
# Probe FPS (with timeout safety ideally, but simple here)
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
if cap.isOpened():
|
||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
cap.release()
|
||||
except Exception:
|
||||
logger.warning("Could not probe FPS from source.")
|
||||
|
||||
self.stopped = False
|
||||
self.thread = threading.Thread(target=self._update, daemon=True)
|
||||
self._start()
|
||||
logger.info("Initialized ThreadedRTSPSource for %s", source)
|
||||
|
||||
@property
|
||||
def fps(self) -> float:
|
||||
return self._fps
|
||||
|
||||
def _start(self):
|
||||
self.thread.start()
|
||||
|
||||
def _update(self):
|
||||
logger.info("Reader thread started for %s", self.source)
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
frame_id = 0
|
||||
while not self.stopped:
|
||||
if not cap.isOpened():
|
||||
# Try reconnecting
|
||||
logger.warning("Stream not open, retrying...")
|
||||
time.sleep(1)
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
continue
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
# Stream ended or failed
|
||||
logger.warning("Failed to read frame from RTSP, retrying...")
|
||||
time.sleep(0.1)
|
||||
continue
|
||||
|
||||
# If queue is full, remove the oldest frame to keep latest
|
||||
if self.queue.full():
|
||||
try:
|
||||
self.queue.get_nowait()
|
||||
except Empty:
|
||||
pass
|
||||
|
||||
self.queue.put(Frame(timestamp=time.time(), data=frame, frame_id=frame_id))
|
||||
frame_id += 1
|
||||
cap.release()
|
||||
logger.info("RTSP Thread stopped.")
|
||||
|
||||
def stream(self) -> Iterator[Frame]:
|
||||
while not self.stopped:
|
||||
try:
|
||||
yield self.queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
|
||||
def stop(self):
|
||||
self.stopped = True
|
||||
self.thread.join()
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Stream sources for the pipeline."""
|
||||
|
||||
from pipeline.sources.file import FileStreamSource
|
||||
from pipeline.sources.stream import ThreadedStreamSource
|
||||
|
||||
# Alias for backwards compatibility
|
||||
ThreadedRTSPSource = ThreadedStreamSource
|
||||
|
||||
__all__ = [
|
||||
"FileStreamSource",
|
||||
"ThreadedStreamSource",
|
||||
"ThreadedRTSPSource", # Alias
|
||||
]
|
||||
@@ -0,0 +1,63 @@
|
||||
"""File-based video source."""
|
||||
|
||||
import glob
|
||||
import logging
|
||||
import time
|
||||
from typing import Iterator
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.protocols import StreamSource
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileStreamSource(StreamSource):
|
||||
"""Reads frames from a video file."""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = path
|
||||
self._fps = 30.0
|
||||
|
||||
# Probe for FPS if it's a single file
|
||||
if "*" not in self.path:
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.path)
|
||||
if cap.isOpened():
|
||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
cap.release()
|
||||
except Exception as e:
|
||||
logger.warning("Could not probe FPS from %s: %s", path, e)
|
||||
|
||||
@property
|
||||
def fps(self) -> float:
|
||||
return self._fps
|
||||
|
||||
def stream(self) -> Iterator[Frame]:
|
||||
# Handle glob patterns for image sequences
|
||||
if "*" in self.path:
|
||||
files = sorted(glob.glob(self.path))
|
||||
logger.info("Loading %d images from pattern: %s", len(files), self.path)
|
||||
for frame_id, file in enumerate(files):
|
||||
img = cv2.imread(file)
|
||||
if img is not None:
|
||||
yield Frame(timestamp=time.time(), data=img, frame_id=frame_id)
|
||||
return
|
||||
|
||||
logger.info("Opening video file: %s", self.path)
|
||||
cap = cv2.VideoCapture(self.path)
|
||||
frame_id = 0
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
logger.info("End of video file reached.")
|
||||
break
|
||||
|
||||
yield Frame(timestamp=time.time(), data=frame, frame_id=frame_id)
|
||||
frame_id += 1
|
||||
cap.release()
|
||||
|
||||
def stop(self):
|
||||
"""No-op for file source (stream ends naturally)."""
|
||||
pass
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Threaded source for RTSP streams and webcams."""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from queue import Empty, Queue
|
||||
from typing import Iterator
|
||||
|
||||
import cv2
|
||||
|
||||
from pipeline.protocols import StreamSource
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ThreadedStreamSource(StreamSource):
|
||||
"""
|
||||
Reads from RTSP stream or webcam in a background thread.
|
||||
|
||||
Prevents buffer buildup and allows main processing loop to run
|
||||
at its own pace without stalling the decoding.
|
||||
"""
|
||||
|
||||
def __init__(self, source: str | int, queue_size: int = 5):
|
||||
self.source = source
|
||||
self.queue: Queue[Frame] = Queue(maxsize=queue_size)
|
||||
self._fps = 30.0
|
||||
|
||||
# Probe FPS
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
if cap.isOpened():
|
||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
cap.release()
|
||||
except cv2.error: # pylint: disable=catching-non-exception
|
||||
logger.warning("Could not probe FPS from source.")
|
||||
|
||||
self.stopped = False
|
||||
self.thread = threading.Thread(target=self._reader_loop, daemon=True)
|
||||
self.thread.start()
|
||||
logger.info("Initialized ThreadedStreamSource for %s", source)
|
||||
|
||||
@property
|
||||
def fps(self) -> float:
|
||||
return self._fps
|
||||
|
||||
def _reader_loop(self):
|
||||
"""Background thread that continuously reads frames."""
|
||||
logger.info("Reader thread started for %s", self.source)
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
frame_id = 0
|
||||
|
||||
while not self.stopped:
|
||||
if not cap.isOpened():
|
||||
logger.warning("Stream not open, retrying...")
|
||||
time.sleep(1)
|
||||
cap = cv2.VideoCapture(self.source)
|
||||
continue
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
logger.warning("Failed to read frame, retrying...")
|
||||
time.sleep(0.1)
|
||||
continue
|
||||
|
||||
# Drop oldest if queue full (keep latest frames)
|
||||
if self.queue.full():
|
||||
try:
|
||||
self.queue.get_nowait()
|
||||
except Empty:
|
||||
pass
|
||||
|
||||
self.queue.put(Frame(timestamp=time.time(), data=frame, frame_id=frame_id))
|
||||
frame_id += 1
|
||||
|
||||
cap.release()
|
||||
logger.info("Reader thread stopped.")
|
||||
|
||||
def stream(self) -> Iterator[Frame]:
|
||||
"""Yield frames from the queue."""
|
||||
while not self.stopped:
|
||||
try:
|
||||
yield self.queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
|
||||
def stop(self):
|
||||
"""Stop the reader thread and release resources."""
|
||||
self.stopped = True
|
||||
if self.thread.is_alive():
|
||||
self.thread.join(timeout=2.0)
|
||||
@@ -58,62 +58,6 @@ class Detection:
|
||||
metadata: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
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 label, dets in by_label.items():
|
||||
result.extend(nms(dets, iou_threshold))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class Ticket:
|
||||
"""Represents a potential alert event."""
|
||||
|
||||
Reference in new issue
Block a user