Cooked I Am

This commit is contained in:
lda committed 2026-01-15 23:21:37 +07:00
1 parent 7b2813be20
commit d57a48eec6
20 files changed
+665 -276

No files matched your search

+1 -1
View File
@@ -5,5 +5,5 @@
"**/*.rpymc": true,
"**/cache/": true
},
"pylint.args": ["--generated-members", "cv2.*"],
"pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"],
}
+10 -2
View File
@@ -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
View File
@@ -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__)
+15
View File
@@ -0,0 +1,15 @@
"""Model wrappers for the pipeline."""
from pipeline.models.yolo import (
DEVICE,
YoloCigaretteDetector,
YoloClassifier,
YoloPersonDetector,
)
__all__ = [
"DEVICE",
"YoloPersonDetector",
"YoloCigaretteDetector",
"YoloClassifier",
]
+16 -18
View File
@@ -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
+41
View File
@@ -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",
]
+117
View File
@@ -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,
)
+73
View File
@@ -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)
+61
View File
@@ -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
+4
View File
@@ -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",
-63
View File
@@ -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,
)
+1 -1
View File
@@ -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__)
+102
View File
@@ -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)
+53
View File
@@ -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)
+1 -1
View File
@@ -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__)
-133
View File
@@ -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()
+13
View File
@@ -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
]
+63
View File
@@ -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
+92
View File
@@ -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)
-56
View File
@@ -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."""