Mystic Dicky
This commit is contained in:
lda committed
1 parent
fa32c7b28c
commit
7b2813be20
9 files changed
+691
-337
No files matched your search
@@ -0,0 +1,202 @@
|
||||
import cv2
|
||||
import time
|
||||
import threading
|
||||
from collections import deque
|
||||
from ultralytics.models import YOLO
|
||||
|
||||
# ========= CẤU HÌNH =========
|
||||
CAM_INDEX = 0
|
||||
BUFFER_SIZE = 10
|
||||
PROCESS_FPS = 2
|
||||
DET_THRESH = 0.8
|
||||
CLS_THRESH = 0.8
|
||||
|
||||
DET_MODEL_PATH = r"F:\Intership\Smoke\test\cig-det-best.pt"
|
||||
CLS_MODEL_PATH = r"F:\Intership\Smoke\test\cig-cls-best.pt"
|
||||
|
||||
# Nếu classifier có 2 lớp: smoke / no_smoke (hoặc smoking / not_smoking)
|
||||
# Bạn đặt tên lớp "hút thuốc" tại đây cho chắc chắn.
|
||||
SMOKING_CLASS_NAME = "smoke" # đổi thành "smoking" nếu model bạn đặt tên như vậy
|
||||
|
||||
# ========= BIẾN CHIA SẺ =========
|
||||
frame_buffer = deque(maxlen=BUFFER_SIZE)
|
||||
buffer_lock = threading.Lock()
|
||||
stop_event = threading.Event()
|
||||
|
||||
# ========= LOAD MODEL =========
|
||||
det_model = YOLO(DET_MODEL_PATH)
|
||||
cls_model = YOLO(CLS_MODEL_PATH)
|
||||
|
||||
|
||||
def capture_and_display():
|
||||
cap = cv2.VideoCapture(CAM_INDEX)
|
||||
if not cap.isOpened():
|
||||
print("Không mở được camera.")
|
||||
stop_event.set()
|
||||
return
|
||||
|
||||
while not stop_event.is_set():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
print("Không đọc được frame từ camera.")
|
||||
break
|
||||
|
||||
with buffer_lock:
|
||||
frame_buffer.append(frame)
|
||||
|
||||
cv2.imshow("Realtime Camera", frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
stop_event.set()
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def get_latest_frame():
|
||||
with buffer_lock:
|
||||
if not frame_buffer:
|
||||
return None
|
||||
return frame_buffer[-1].copy()
|
||||
|
||||
|
||||
def run_detection(frame):
|
||||
"""
|
||||
Return:
|
||||
best_conf: float
|
||||
best_box: (x1,y1,x2,y2) int or None
|
||||
"""
|
||||
results = det_model.predict(source=frame, verbose=False)
|
||||
r = results[0]
|
||||
|
||||
if r.boxes is None or len(r.boxes) == 0:
|
||||
return 0.0, None
|
||||
|
||||
confs = r.boxes.conf.cpu().numpy()
|
||||
boxes = r.boxes.xyxy.cpu().numpy()
|
||||
|
||||
best_idx = confs.argmax()
|
||||
best_conf = float(confs[best_idx])
|
||||
x1, y1, x2, y2 = boxes[best_idx]
|
||||
best_box = (int(x1), int(y1), int(x2), int(y2))
|
||||
|
||||
return best_conf, best_box
|
||||
|
||||
|
||||
def crop_from_box(frame, box, pad=10):
|
||||
"""
|
||||
Crop theo box + padding nhẹ.
|
||||
Nếu box None -> dùng full frame.
|
||||
"""
|
||||
if box is None:
|
||||
return frame
|
||||
|
||||
x1, y1, x2, y2 = box
|
||||
h, w = frame.shape[:2]
|
||||
|
||||
x1 = max(0, x1 - pad)
|
||||
y1 = max(0, y1 - pad)
|
||||
x2 = min(w - 1, x2 + pad)
|
||||
y2 = min(h - 1, y2 + pad)
|
||||
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
return frame
|
||||
return frame[y1:y2, x1:x2]
|
||||
|
||||
|
||||
def run_classification(img):
|
||||
"""
|
||||
Return:
|
||||
cls_name: str
|
||||
cls_conf: float
|
||||
"""
|
||||
results = cls_model.predict(source=img, verbose=False)
|
||||
r = results[0]
|
||||
|
||||
probs = r.probs
|
||||
top1 = int(probs.top1)
|
||||
top1conf = float(probs.top1conf)
|
||||
|
||||
names = r.names
|
||||
if isinstance(names, dict):
|
||||
cls_name = names.get(top1, str(top1))
|
||||
else:
|
||||
cls_name = names[top1] if top1 < len(names) else str(top1)
|
||||
|
||||
return cls_name, top1conf
|
||||
|
||||
|
||||
def is_smoking_by_class(cls_name, cls_conf):
|
||||
"""
|
||||
CHỈ báo hút thuốc khi:
|
||||
- confidence đủ cao
|
||||
- và class đúng là lớp hút thuốc (SMOKING_CLASS_NAME)
|
||||
"""
|
||||
if cls_conf < CLS_THRESH:
|
||||
return False
|
||||
|
||||
# bắt buộc phải biết class "hút thuốc" là gì
|
||||
if not SMOKING_CLASS_NAME:
|
||||
# Nếu bạn không set được tên class, KHÔNG được auto-True
|
||||
return False
|
||||
|
||||
return cls_name.strip().lower() == SMOKING_CLASS_NAME.strip().lower()
|
||||
|
||||
|
||||
def processing_loop():
|
||||
interval = 1.0 / PROCESS_FPS
|
||||
last_time = 0.0
|
||||
|
||||
while not stop_event.is_set():
|
||||
now = time.time()
|
||||
if now - last_time < interval:
|
||||
time.sleep(0.001)
|
||||
continue
|
||||
last_time = now
|
||||
|
||||
frame = get_latest_frame()
|
||||
if frame is None:
|
||||
continue
|
||||
|
||||
# ===== Step 1: Detection =====
|
||||
det_conf, det_box = run_detection(frame)
|
||||
|
||||
if det_conf >= DET_THRESH:
|
||||
print(f"[ALERT] HÚT THUỐC (det_conf={det_conf:.2f})")
|
||||
continue
|
||||
|
||||
# ===== Step 2: det_conf < 0.8 -> crop -> classification =====
|
||||
crop_img = crop_from_box(
|
||||
frame, det_box
|
||||
) # nếu có box thì crop, không có thì full frame
|
||||
cls_name, cls_conf = run_classification(crop_img)
|
||||
|
||||
if is_smoking_by_class(cls_name, cls_conf):
|
||||
print(
|
||||
f"[ALERT] CÓ HÚT THUỐC (cls={cls_name}, cls_conf={cls_conf:.2f}, det_conf={det_conf:.2f})"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"[OK] KHÔNG HÚT/CHƯA CHẮC (cls={cls_name}, cls_conf={cls_conf:.2f}, det_conf={det_conf:.2f})"
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
t1 = threading.Thread(target=capture_and_display, daemon=True)
|
||||
t2 = threading.Thread(target=processing_loop, daemon=True)
|
||||
|
||||
t1.start()
|
||||
t2.start()
|
||||
|
||||
try:
|
||||
while not stop_event.is_set():
|
||||
time.sleep(0.1)
|
||||
except KeyboardInterrupt:
|
||||
stop_event.set()
|
||||
|
||||
t1.join()
|
||||
t2.join()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,7 +6,12 @@ from pathlib import Path
|
||||
|
||||
from pipeline.core import SmokingDetectionPipeline
|
||||
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
|
||||
from pipeline.outputs import FileAlertStore, SmartEventSink, VideoFileStore
|
||||
from pipeline.sinks import (
|
||||
FileAlertStore,
|
||||
SmartEventSink,
|
||||
VideoFileStore,
|
||||
WindowDisplaySink,
|
||||
)
|
||||
from pipeline.sources import FileStreamSource, ThreadedRTSPSource
|
||||
|
||||
|
||||
@@ -83,6 +88,14 @@ def main():
|
||||
help="Path to output video file (will also generate .jsonl)",
|
||||
)
|
||||
parser.add_argument("--debug", action="store_true", help="Increase verbosity")
|
||||
parser.add_argument(
|
||||
"--display", action="store_true", help="Show live preview window"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save",
|
||||
action="store_true",
|
||||
help="Don't save output video (useful with --display)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -118,23 +131,38 @@ def main():
|
||||
# Simple file reader
|
||||
source = FileStreamSource(args.source)
|
||||
|
||||
# Initialize Output Sink
|
||||
# Pass FPS from source so output video plays at correct speed
|
||||
video_sink = VideoFileStore(args.output, fps=int(source.fps))
|
||||
# Build sinks list based on arguments
|
||||
sinks = []
|
||||
video_sink = None
|
||||
smart_sink = None
|
||||
alert_store = None
|
||||
display_sink = None
|
||||
|
||||
# Initialize Smart Sink (Optional, writes small clips only when smoking is detected)
|
||||
smart_sink = SmartEventSink("events", fps=source.fps)
|
||||
if not args.no_save:
|
||||
# Initialize Output Sink
|
||||
# Pass FPS from source so output video plays at correct speed
|
||||
video_sink = VideoFileStore(args.output, fps=int(source.fps))
|
||||
sinks.append(video_sink)
|
||||
|
||||
# Initialize Alert Store
|
||||
# This persists confirmed smoking events to a dedicated log file
|
||||
# Initialize Smart Sink (writes small clips only when smoking is detected)
|
||||
smart_sink = SmartEventSink("events", fps=source.fps)
|
||||
sinks.append(smart_sink)
|
||||
|
||||
# Initialize Alert Store (always active)
|
||||
alert_store = FileAlertStore("alerts.jsonl")
|
||||
sinks.append(alert_store)
|
||||
|
||||
# Live display window
|
||||
if args.display:
|
||||
display_sink = WindowDisplaySink("Smoking Detection", scale=1.0)
|
||||
sinks.append(display_sink)
|
||||
|
||||
pipeline = SmokingDetectionPipeline(
|
||||
person_model=person_model,
|
||||
cigarette_model=cigarette_model,
|
||||
verifier_model=verifier,
|
||||
vlm_model=None, # Optional
|
||||
sinks=[video_sink, alert_store, smart_sink],
|
||||
sinks=sinks,
|
||||
inference_interval=0.5, # Slow mode (2 FPS) when nobody is home
|
||||
active_interval=0.0, # Max speed when person detected
|
||||
cooldown_frames=int(
|
||||
@@ -142,23 +170,18 @@ def main():
|
||||
), # Stay active for 2 seconds after person leaves
|
||||
)
|
||||
|
||||
print("Starting pipeline...")
|
||||
logger.info("Starting pipeline with %d sinks...", len(sinks))
|
||||
try:
|
||||
# Pass the already initialized source
|
||||
pipeline.process_stream(source)
|
||||
except KeyboardInterrupt:
|
||||
print("Stopped by user.")
|
||||
finally:
|
||||
if args.rtsp and hasattr(source, "stop"):
|
||||
source.stop() # pyright: ignore[reportAttributeAccessIssue]
|
||||
source.stop()
|
||||
|
||||
# Explicitly close all sinks
|
||||
if video_sink:
|
||||
video_sink.close()
|
||||
if smart_sink:
|
||||
smart_sink.close()
|
||||
if alert_store:
|
||||
alert_store.close()
|
||||
# Close all sinks
|
||||
for sink in sinks:
|
||||
sink.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -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