Mystic Dicky
This commit is contained in:
lda committed
1 parent
fa32c7b28c
commit
7b2813be20
9 files changed
+691
-337
No files matched your search
@@ -1,318 +0,0 @@
|
||||
import collections
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from queue import Empty, Queue
|
||||
from typing import Deque, List, Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from pipeline.protocols import Sink
|
||||
from pipeline.types import Detection, 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=2,
|
||||
post_event_seconds=3,
|
||||
):
|
||||
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 = 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.")
|
||||
|
||||
|
||||
class FileAlertStore:
|
||||
"""
|
||||
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(
|
||||
f"[ALERT SAVED] {ticket.confidence:.2f} confirmed at {ticket.stage_confirmed}"
|
||||
)
|
||||
|
||||
def close(self):
|
||||
self.file.close()
|
||||
|
||||
|
||||
class VideoFileStore:
|
||||
"""
|
||||
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(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")
|
||||
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 (30 frames ~1 second) 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
|
||||
self._draw_detections(vis_image, frame.detections)
|
||||
self._draw_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 _draw_detections(self, image: np.ndarray, detections: List[Detection]):
|
||||
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_overlay(self, image: np.ndarray, frame_id: int):
|
||||
cv2.putText(
|
||||
image,
|
||||
f"Frame: {frame_id}",
|
||||
(10, 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
1.0,
|
||||
(0, 255, 0),
|
||||
2,
|
||||
)
|
||||
|
||||
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.")
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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.")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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.")
|
||||
@@ -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.")
|
||||
@@ -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.")
|
||||
Reference in new issue
Block a user