Mystic Dicky

pipeline/sinks/__init__.py

@@ -0,0 +1,13 @@
+"""Sinks for pipeline output handling."""
+
+from pipeline.sinks.alerts import FileAlertStore
+from pipeline.sinks.display import WindowDisplaySink
+from pipeline.sinks.events import SmartEventSink
+from pipeline.sinks.video import VideoFileStore
+
+__all__ = [
+    "FileAlertStore",
+    "SmartEventSink",
+    "VideoFileStore",
+    "WindowDisplaySink",
+]

pipeline/sinks/alerts.py

@@ -0,0 +1,48 @@
+"""Alert persistence sink - writes tickets to JSONL."""
+
+import json
+import logging
+
+from pipeline.protocols import Sink
+from pipeline.types import Frame
+
+logger = logging.getLogger(__name__)
+
+
+class FileAlertStore(Sink):
+    """
+    Persists actionable alerts (Tickets) to a JSONL file.
+    Use this for integrating with backend systems or audit logs.
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+        self.file = open(path, "a", encoding="utf-8")  # Append mode
+        logger.info("Initialized FileAlertStore writing to %s", path)
+
+    def handle(self, frame: Frame):
+        if not frame.ticket:
+            return
+
+        ticket = frame.ticket
+        record = {
+            "type": "alert",
+            "frame_id": ticket.frame_id,
+            "confidence": ticket.confidence,
+            "stage": ticket.stage_confirmed,
+            "metadata": ticket.metadata,
+        }
+        # Note: We skip serializing 'evidence' (raw arrays) to JSON
+        # In a real app, you'd save the image to disk and link the path here.
+
+        self.file.write(json.dumps(record) + "\n")
+        self.file.flush()
+        logger.info(
+            "[ALERT SAVED] %.2f confirmed at %s",
+            ticket.confidence,
+            ticket.stage_confirmed,
+        )
+
+    def close(self):
+        self.file.close()
+        logger.info("FileAlertStore closed.")

pipeline/sinks/base.py

@@ -0,0 +1,63 @@
+"""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,
+    )

pipeline/sinks/display.py

@@ -0,0 +1,64 @@
+"""Live window display sink for real-time visualization."""
+
+import logging
+
+import cv2
+
+from pipeline.protocols import Sink
+from pipeline.sinks.base import draw_detections, draw_frame_overlay
+from pipeline.types import Frame
+
+logger = logging.getLogger(__name__)
+
+
+class WindowDisplaySink(Sink):
+    """
+    Displays frames in a live OpenCV window.
+    Press 'q' to signal quit (sets self.quit_requested).
+    """
+
+    def __init__(self, window_name: str = "Smoking Detection", scale: float = 1.0):
+        self.window_name = window_name
+        self.scale = scale
+        self.quit_requested = False
+        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
+        logger.info("Initialized WindowDisplaySink: %s", window_name)
+
+    def handle(self, frame: Frame):
+        if self.quit_requested:
+            return
+
+        vis = frame.data.copy()
+
+        # Draw detections (same style as VideoFileStore)
+        draw_detections(vis, frame.detections)
+        draw_frame_overlay(vis, frame.frame_id)
+
+        # Alert indicator
+        if frame.ticket:
+            cv2.rectangle(vis, (0, 0), (vis.shape[1], 40), (0, 0, 255), -1)
+            cv2.putText(
+                vis,
+                f"SMOKING DETECTED ({frame.ticket.confidence:.2f})",
+                (10, 30),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (255, 255, 255),
+                2,
+            )
+
+        # Scale if needed
+        if self.scale != 1.0:
+            vis = cv2.resize(vis, None, fx=self.scale, fy=self.scale)
+
+        cv2.imshow(self.window_name, vis)
+
+        # Check for quit key
+        key = cv2.waitKey(1) & 0xFF
+        if key == ord("q"):
+            self.quit_requested = True
+            logger.info("Quit requested via 'q' key")
+
+    def close(self):
+        cv2.destroyWindow(self.window_name)
+        logger.info("WindowDisplaySink closed.")

pipeline/sinks/events.py

@@ -0,0 +1,135 @@
+"""Smart event clip recorder - captures context around detections."""
+
+import collections
+import logging
+import time
+from pathlib import Path
+from typing import Deque, Optional
+
+import cv2
+
+from pipeline.protocols import Sink
+from pipeline.types import Frame
+
+logger = logging.getLogger(__name__)
+
+
+class SmartEventSink(Sink):
+    """
+    Smart Event Sink:
+    - Maintains a ring buffer of recent frames (e.g., 2 seconds context).
+    - When an alert (Ticket) occurs, it triggers a 'recording' state.
+    - Captures pre-event context + event + post-event context.
+    - Saves these 'events' as separate video clips or sequences.
+    """
+
+    def __init__(
+        self,
+        output_dir: str,
+        fps: float = 30.0,
+        pre_event_seconds: float = 2.0,
+        post_event_seconds: float = 3.0,
+    ):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.fps = fps
+        self.buffer_size = int(fps * pre_event_seconds)
+        self.post_event_frames = int(fps * post_event_seconds)
+
+        self.buffer: Deque[Frame] = collections.deque(maxlen=self.buffer_size)
+
+        # State
+        self.recording = False
+        self.remaining_post_frames = 0
+        self.current_event_writer: Optional[cv2.VideoWriter] = None
+        self.current_event_id: Optional[int] = None
+
+        logger.info("Initialized SmartEventSink in %s", output_dir)
+
+    def handle(self, frame: Frame):
+        # Always add to buffer if not recording to keep context
+        if not self.recording:
+            self.buffer.append(frame)
+
+        # Check for trigger
+        if frame.ticket:
+            self._trigger_event(frame)
+
+        # If recording, write frame
+        if self.recording:
+            self._write_to_event(frame)
+
+            # Decrement counter
+            # If we keep getting tickets, we keep updating the counter (extending recording)
+            if frame.ticket:
+                if self.remaining_post_frames != self.post_event_frames:
+                    logger.debug(
+                        "Event %s extended due to new detection.", self.current_event_id
+                    )
+                self.remaining_post_frames = self.post_event_frames
+            else:
+                self.remaining_post_frames -= 1
+
+            if self.remaining_post_frames <= 0:
+                self._stop_event()
+
+    def _trigger_event(self, trigger_frame: Frame):
+        if not self.recording:
+            # Start new event
+            self.recording = True
+            self.current_event_id = int(time.time())
+            self.remaining_post_frames = self.post_event_frames
+
+            filename = self.output_dir / f"event_{self.current_event_id}.mp4"
+            h, w = trigger_frame.data.shape[:2]
+
+            self.current_event_writer = cv2.VideoWriter(
+                str(filename), cv2.VideoWriter.fourcc(*"mp4v"), self.fps, (w, h)
+            )
+
+            # Dump buffer (Pre-event context)
+            while self.buffer:
+                f = self.buffer.popleft()
+                self._write_to_event(f)
+
+            logger.info("Started recording event %s", self.current_event_id)
+
+    def _write_to_event(self, frame: Frame):
+        if self.current_event_writer:
+            # Draw simple visualization for the clip
+            vis = frame.data.copy()
+            if frame.ticket:
+                cv2.rectangle(vis, (0, 0), (vis.shape[1], 40), (0, 0, 255), -1)
+                cv2.putText(
+                    vis,
+                    "SMOKING DETECTED",
+                    (10, 30),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1,
+                    (255, 255, 255),
+                    2,
+                )
+
+            for det in frame.detections:
+                cv2.rectangle(
+                    vis,
+                    (det.box.x1, det.box.y1),
+                    (det.box.x2, det.box.y2),
+                    (255, 0, 0),
+                    2,
+                )
+
+            self.current_event_writer.write(vis)
+
+    def _stop_event(self):
+        self.recording = False
+        if self.current_event_writer:
+            self.current_event_writer.release()
+            self.current_event_writer = None
+        logger.info("Finished recording event %s", self.current_event_id)
+
+    def close(self):
+        """Releases resources and stops any active recording."""
+        if self.recording:
+            self._stop_event()
+        logger.info("SmartEventSink closed.")

pipeline/sinks/video.py

@@ -0,0 +1,124 @@
+"""Video file output sink with async writing."""
+
+import json
+import logging
+import threading
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Optional
+
+import cv2
+
+from pipeline.protocols import Sink
+from pipeline.sinks.base import draw_detections, draw_frame_overlay
+from pipeline.types import Frame
+
+logger = logging.getLogger(__name__)
+
+
+class VideoFileStore(Sink):
+    """
+    Writes processed frames to a video file with visualization overlays.
+    Also optionally writes metadata sidecar files.
+    Async implementation to prevent blocking the main pipeline.
+    """
+
+    def __init__(
+        self,
+        output_path: str,
+        fps: float = 30.0,
+        fourcc: str = "mp4v",
+        queue_size: int = 60,
+    ):
+        self.output_path = output_path
+        self.metadata_path = str(Path(output_path).with_suffix(".jsonl"))
+        self.fps = fps
+        self.fourcc = fourcc
+        self.writer: Optional[cv2.VideoWriter] = None
+        self.frame_size = None
+
+        # Async setup
+        self.queue: Queue[Frame] = Queue(maxsize=queue_size)
+        self.running = True
+        self.thread = threading.Thread(target=self._writer_loop, daemon=True)
+
+        # Open metadata file
+        self.meta_file = open(self.metadata_path, "w", encoding="utf-8")
+        logger.info("Initialized Async VideoFileStore writing to %s", output_path)
+        self.thread.start()
+
+    def handle(self, frame: Frame):
+        """Enqueue frame for processing."""
+        if not self.running:
+            return
+
+        # We BLOCK here if queue is full to guarantee no frames are dropped.
+        # This may cause the main loop to slow down, but output integrity > speed.
+        # The queue size (60 frames ~2 seconds) provides a buffer for burst writes.
+        self.queue.put(frame, block=True)
+
+    def _writer_loop(self):
+        while self.running or not self.queue.empty():
+            try:
+                frame = self.queue.get(timeout=1.0)
+            except Empty:
+                continue
+
+            try:
+                self._process_frame(frame)
+            except Exception as e:
+                logger.error("Error processing video output: %s", e)
+            finally:
+                self.queue.task_done()
+
+    def _process_frame(self, frame: Frame):
+        # Heavy lifting happens here in the thread
+        vis_image = frame.data.copy()
+
+        # 1. Visualization
+        draw_detections(vis_image, frame.detections)
+        draw_frame_overlay(vis_image, frame.frame_id)
+
+        # 2. Write Video
+        if self.writer is None:
+            h, w = vis_image.shape[:2]
+            self.frame_size = (w, h)
+            self.writer = cv2.VideoWriter(
+                self.output_path,
+                cv2.VideoWriter.fourcc(*self.fourcc),
+                self.fps,
+                self.frame_size,
+            )
+
+        self.writer.write(vis_image)
+
+        # 3. Write Metadata
+        self._write_metadata(frame)
+
+    def _write_metadata(self, frame: Frame):
+        meta = {
+            "frame_id": frame.frame_id,
+            "timestamp": frame.timestamp,
+            "detections": [
+                {
+                    "label": d.label,
+                    "score": d.score,
+                    "box": [d.box.x1, d.box.y1, d.box.x2, d.box.y2],
+                    "meta": d.metadata,
+                }
+                for d in frame.detections
+            ],
+        }
+        self.meta_file.write(json.dumps(meta) + "\n")
+        self.meta_file.flush()
+
+    def close(self):
+        self.running = False
+        if self.thread.is_alive():
+            self.thread.join(timeout=2.0)
+
+        if self.writer:
+            self.writer.release()
+        if self.meta_file:
+            self.meta_file.close()
+        logger.info("VideoFileStore closed.")