Compare commits
2
Commits
No files matched your search
Vendored
+1
-1
@@ -5,5 +5,5 @@
|
|||||||
"**/*.rpymc": true,
|
"**/*.rpymc": true,
|
||||||
"**/cache/": true
|
"**/cache/": true
|
||||||
},
|
},
|
||||||
"pylint.args": ["--generated-members", "cv2.*"],
|
"pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"],
|
||||||
}
|
}
|
||||||
@@ -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,14 @@ from pathlib import Path
|
|||||||
|
|
||||||
from pipeline.core import SmokingDetectionPipeline
|
from pipeline.core import SmokingDetectionPipeline
|
||||||
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
|
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
|
||||||
from pipeline.outputs import FileAlertStore, SmartEventSink, VideoFileStore
|
from pipeline.sinks import (
|
||||||
|
FileAlertStore,
|
||||||
|
MetricsSink,
|
||||||
|
ScreenshotSink,
|
||||||
|
SmartEventSink,
|
||||||
|
VideoFileStore,
|
||||||
|
WindowDisplaySink,
|
||||||
|
)
|
||||||
from pipeline.sources import FileStreamSource, ThreadedRTSPSource
|
from pipeline.sources import FileStreamSource, ThreadedRTSPSource
|
||||||
|
|
||||||
|
|
||||||
@@ -83,6 +90,14 @@ def main():
|
|||||||
help="Path to output video file (will also generate .jsonl)",
|
help="Path to output video file (will also generate .jsonl)",
|
||||||
)
|
)
|
||||||
parser.add_argument("--debug", action="store_true", help="Increase verbosity")
|
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()
|
args = parser.parse_args()
|
||||||
|
|
||||||
@@ -118,23 +133,44 @@ def main():
|
|||||||
# Simple file reader
|
# Simple file reader
|
||||||
source = FileStreamSource(args.source)
|
source = FileStreamSource(args.source)
|
||||||
|
|
||||||
# Initialize Output Sink
|
# Build sinks list based on arguments
|
||||||
# Pass FPS from source so output video plays at correct speed
|
sinks = []
|
||||||
video_sink = VideoFileStore(args.output, fps=int(source.fps))
|
video_sink = None
|
||||||
|
smart_sink = None
|
||||||
|
|
||||||
# Initialize Smart Sink (Optional, writes small clips only when smoking is detected)
|
# Metrics sink (always active - provides performance summary)
|
||||||
smart_sink = SmartEventSink("events", fps=source.fps)
|
metrics_sink = MetricsSink(log_interval=100)
|
||||||
|
sinks.append(metrics_sink)
|
||||||
|
|
||||||
# Initialize Alert Store
|
if not args.no_save:
|
||||||
# This persists confirmed smoking events to a dedicated log file
|
# 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 Smart Sink (writes small clips only when smoking is detected)
|
||||||
|
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")
|
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(
|
pipeline = SmokingDetectionPipeline(
|
||||||
person_model=person_model,
|
person_model=person_model,
|
||||||
cigarette_model=cigarette_model,
|
cigarette_model=cigarette_model,
|
||||||
verifier_model=verifier,
|
verifier_model=verifier,
|
||||||
vlm_model=None, # Optional
|
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
|
inference_interval=0.5, # Slow mode (2 FPS) when nobody is home
|
||||||
active_interval=0.0, # Max speed when person detected
|
active_interval=0.0, # Max speed when person detected
|
||||||
cooldown_frames=int(
|
cooldown_frames=int(
|
||||||
@@ -142,23 +178,18 @@ def main():
|
|||||||
), # Stay active for 2 seconds after person leaves
|
), # Stay active for 2 seconds after person leaves
|
||||||
)
|
)
|
||||||
|
|
||||||
print("Starting pipeline...")
|
logger.info("Starting pipeline with %d sinks...", len(sinks))
|
||||||
try:
|
try:
|
||||||
# Pass the already initialized source
|
# Pass the already initialized source
|
||||||
pipeline.process_stream(source)
|
pipeline.process_stream(source)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print("Stopped by user.")
|
print("Stopped by user.")
|
||||||
finally:
|
finally:
|
||||||
if args.rtsp and hasattr(source, "stop"):
|
source.stop()
|
||||||
source.stop() # pyright: ignore[reportAttributeAccessIssue]
|
|
||||||
|
|
||||||
# Explicitly close all sinks
|
# Close all sinks
|
||||||
if video_sink:
|
for sink in sinks:
|
||||||
video_sink.close()
|
sink.close()
|
||||||
if smart_sink:
|
|
||||||
smart_sink.close()
|
|
||||||
if alert_store:
|
|
||||||
alert_store.close()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
+2
-1
@@ -5,8 +5,9 @@ from collections import deque
|
|||||||
from concurrent.futures import Future
|
from concurrent.futures import Future
|
||||||
from typing import List, Optional, Tuple
|
from typing import List, Optional, Tuple
|
||||||
|
|
||||||
|
from pipeline.ops import nms_by_class
|
||||||
from pipeline.protocols import Model, Sink, StreamSource, Verifier
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,15 @@
|
|||||||
|
"""Model wrappers for the pipeline."""
|
||||||
|
|
||||||
|
from pipeline.models.yolo import (
|
||||||
|
DEVICE,
|
||||||
|
YoloCigaretteDetector,
|
||||||
|
YoloClassifier,
|
||||||
|
YoloPersonDetector,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"DEVICE",
|
||||||
|
"YoloPersonDetector",
|
||||||
|
"YoloCigaretteDetector",
|
||||||
|
"YoloClassifier",
|
||||||
|
]
|
||||||
@@ -1,3 +1,5 @@
|
|||||||
|
"""YOLO model wrappers for detection and classification."""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from typing import Any, List
|
from typing import Any, List
|
||||||
|
|
||||||
@@ -18,12 +20,13 @@ else:
|
|||||||
|
|
||||||
|
|
||||||
class YoloPersonDetector(Model):
|
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)
|
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]:
|
def predict(self, input_data: Any) -> List[Detection]:
|
||||||
# input_data: numpy array (H, W, C)
|
|
||||||
results = self.model(
|
results = self.model(
|
||||||
input_data, classes=[0], verbose=False, device=DEVICE
|
input_data, classes=[0], verbose=False, device=DEVICE
|
||||||
) # 0 is person
|
) # 0 is person
|
||||||
@@ -45,13 +48,14 @@ class YoloPersonDetector(Model):
|
|||||||
|
|
||||||
|
|
||||||
class YoloCigaretteDetector(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)
|
logger.info("Loading Cigarette Detector: %s", model_path)
|
||||||
self.model = YOLO(model_path)
|
self.model = YOLO(model_path)
|
||||||
|
self.conf_threshold = confidence_threshold
|
||||||
|
|
||||||
def predict(self, input_data: Any) -> List[Detection]:
|
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)
|
results = self.model(input_data, verbose=False, device=DEVICE)
|
||||||
|
|
||||||
detections = []
|
detections = []
|
||||||
@@ -60,8 +64,7 @@ class YoloCigaretteDetector(Model):
|
|||||||
for box in boxes:
|
for box in boxes:
|
||||||
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
||||||
conf = float(box.conf[0])
|
conf = float(box.conf[0])
|
||||||
# Maybe filter by conf here if needed
|
if conf > self.conf_threshold:
|
||||||
if conf > 0.2:
|
|
||||||
detections.append(
|
detections.append(
|
||||||
Detection(
|
Detection(
|
||||||
box=BoundingBox(x1, y1, x2, y2, conf),
|
box=BoundingBox(x1, y1, x2, y2, conf),
|
||||||
@@ -73,18 +76,18 @@ class YoloCigaretteDetector(Model):
|
|||||||
|
|
||||||
|
|
||||||
class YoloClassifier(Verifier):
|
class YoloClassifier(Verifier):
|
||||||
|
"""Classifies smoking using YOLO classification model."""
|
||||||
|
|
||||||
def __init__(self, model_path: str):
|
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)
|
self.model = YOLO(model_path)
|
||||||
|
|
||||||
def verify(self, input_data: Any) -> float:
|
def verify(self, input_data: Any) -> float:
|
||||||
# input_data: numpy array (H, W, C) - Usually a crop
|
|
||||||
if input_data.size == 0:
|
if input_data.size == 0:
|
||||||
return 0.0
|
return 0.0
|
||||||
|
|
||||||
results = self.model(input_data, verbose=False, device=DEVICE)
|
results = self.model(input_data, verbose=False, device=DEVICE)
|
||||||
|
|
||||||
# Check if we got any results
|
|
||||||
if not results:
|
if not results:
|
||||||
return 0.0
|
return 0.0
|
||||||
|
|
||||||
@@ -92,16 +95,11 @@ class YoloClassifier(Verifier):
|
|||||||
if probs is None:
|
if probs is None:
|
||||||
return 0.0
|
return 0.0
|
||||||
|
|
||||||
# User specified Class 0 is "smoke"
|
# Class 0 is "smoke" - return its probability
|
||||||
# We return the probability of class 0 specifically.
|
|
||||||
# Ultralytics probs.data is a tensor of scores for each class.
|
|
||||||
try:
|
try:
|
||||||
return float(probs.data[0])
|
return float(probs.data[0])
|
||||||
except (AttributeError, IndexError):
|
except (AttributeError, IndexError):
|
||||||
# Fallback if structure is unexpected, though data[0] is standard for classification
|
logger.warning("Could not retrieve Class 0 probability. Checking top1.")
|
||||||
logger.warning(
|
|
||||||
"Could not retrieve Class 0 probability. Defaulting to top1conf if class 0."
|
|
||||||
)
|
|
||||||
if hasattr(probs, "top1") and probs.top1 == 0:
|
if hasattr(probs, "top1") and probs.top1 == 0:
|
||||||
return float(probs.top1conf)
|
return float(probs.top1conf)
|
||||||
return 0.0
|
return 0.0
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
"""Operations and utility functions."""
|
||||||
|
|
||||||
|
from pipeline.ops.drawing import (
|
||||||
|
draw_alert_banner,
|
||||||
|
draw_box,
|
||||||
|
draw_detections,
|
||||||
|
draw_fps,
|
||||||
|
draw_frame_overlay,
|
||||||
|
draw_label,
|
||||||
|
draw_zone,
|
||||||
|
get_color,
|
||||||
|
)
|
||||||
|
from pipeline.ops.geometry import (
|
||||||
|
box_center,
|
||||||
|
box_in_zone,
|
||||||
|
expand_box,
|
||||||
|
point_in_box,
|
||||||
|
point_in_polygon,
|
||||||
|
)
|
||||||
|
from pipeline.ops.nms import nms, nms_by_class
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
# NMS
|
||||||
|
"nms",
|
||||||
|
"nms_by_class",
|
||||||
|
# Geometry
|
||||||
|
"expand_box",
|
||||||
|
"point_in_box",
|
||||||
|
"box_center",
|
||||||
|
"point_in_polygon",
|
||||||
|
"box_in_zone",
|
||||||
|
# Drawing
|
||||||
|
"get_color",
|
||||||
|
"draw_box",
|
||||||
|
"draw_label",
|
||||||
|
"draw_detections",
|
||||||
|
"draw_frame_overlay",
|
||||||
|
"draw_alert_banner",
|
||||||
|
"draw_zone",
|
||||||
|
"draw_fps",
|
||||||
|
]
|
||||||
@@ -0,0 +1,117 @@
|
|||||||
|
"""Drawing utilities for visualization."""
|
||||||
|
|
||||||
|
from typing import List, Tuple
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from pipeline.types import Detection
|
||||||
|
|
||||||
|
# Color palette (BGR)
|
||||||
|
COLORS = {
|
||||||
|
"person": (255, 0, 0), # Blue
|
||||||
|
"candidate": (0, 255, 255), # Yellow
|
||||||
|
"verified": (0, 0, 255), # Red
|
||||||
|
"default": (255, 255, 255), # White
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def get_color(label: str) -> Tuple[int, int, int]:
|
||||||
|
"""Get color for a detection label."""
|
||||||
|
for key, color in COLORS.items():
|
||||||
|
if key in label:
|
||||||
|
return color
|
||||||
|
return COLORS["default"]
|
||||||
|
|
||||||
|
|
||||||
|
def draw_box(
|
||||||
|
image: np.ndarray,
|
||||||
|
box: Tuple[int, int, int, int],
|
||||||
|
color: Tuple[int, int, int],
|
||||||
|
thickness: int = 2,
|
||||||
|
):
|
||||||
|
"""Draw a rectangle on image."""
|
||||||
|
x1, y1, x2, y2 = box
|
||||||
|
cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_label(
|
||||||
|
image: np.ndarray,
|
||||||
|
text: str,
|
||||||
|
position: Tuple[int, int],
|
||||||
|
color: Tuple[int, int, int],
|
||||||
|
font_scale: float = 0.5,
|
||||||
|
):
|
||||||
|
"""Draw a label with background."""
|
||||||
|
x, y = position
|
||||||
|
(tw, _), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1)
|
||||||
|
cv2.rectangle(image, (x, y - 20), (x + tw, y), color, -1)
|
||||||
|
cv2.putText(
|
||||||
|
image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), 1
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_detections(image: np.ndarray, detections: List[Detection]):
|
||||||
|
"""Draw all detections with boxes and labels."""
|
||||||
|
for det in detections:
|
||||||
|
color = get_color(det.label)
|
||||||
|
draw_box(image, (det.box.x1, det.box.y1, det.box.x2, det.box.y2), color)
|
||||||
|
draw_label(
|
||||||
|
image, f"{det.label} {det.score:.2f}", (det.box.x1, det.box.y1), color
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_frame_overlay(image: np.ndarray, frame_id: int):
|
||||||
|
"""Draw frame counter overlay."""
|
||||||
|
cv2.putText(
|
||||||
|
image,
|
||||||
|
f"Frame: {frame_id}",
|
||||||
|
(10, 30),
|
||||||
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
|
1.0,
|
||||||
|
(0, 255, 0),
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_alert_banner(image: np.ndarray, text: str = "SMOKING DETECTED"):
|
||||||
|
"""Draw a red alert banner at the top."""
|
||||||
|
cv2.rectangle(image, (0, 0), (image.shape[1], 40), (0, 0, 255), -1)
|
||||||
|
cv2.putText(
|
||||||
|
image,
|
||||||
|
text,
|
||||||
|
(10, 30),
|
||||||
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
|
1,
|
||||||
|
(255, 255, 255),
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_zone(
|
||||||
|
image: np.ndarray,
|
||||||
|
zone: List[Tuple[int, int]],
|
||||||
|
color: Tuple[int, int, int] = (0, 255, 0),
|
||||||
|
alpha: float = 0.3,
|
||||||
|
):
|
||||||
|
"""Draw a semi-transparent polygon zone."""
|
||||||
|
overlay = image.copy()
|
||||||
|
pts = np.array(zone, np.int32).reshape((-1, 1, 2))
|
||||||
|
cv2.fillPoly(overlay, [pts], color)
|
||||||
|
cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
|
||||||
|
cv2.polylines(image, [pts], True, color, 2)
|
||||||
|
|
||||||
|
|
||||||
|
def draw_fps(image: np.ndarray, fps: float):
|
||||||
|
"""Draw FPS counter in top-right corner."""
|
||||||
|
h, w = image.shape[:2]
|
||||||
|
text = f"FPS: {fps:.1f}"
|
||||||
|
cv2.putText(
|
||||||
|
image,
|
||||||
|
text,
|
||||||
|
(w - 120, 30),
|
||||||
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
|
0.7,
|
||||||
|
(0, 255, 0),
|
||||||
|
2,
|
||||||
|
)
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
"""Geometry utilities for bounding box operations."""
|
||||||
|
|
||||||
|
from typing import List, Tuple
|
||||||
|
|
||||||
|
from pipeline.types import BoundingBox
|
||||||
|
|
||||||
|
|
||||||
|
def expand_box(
|
||||||
|
box: BoundingBox,
|
||||||
|
margin: float,
|
||||||
|
image_size: Tuple[int, int],
|
||||||
|
) -> Tuple[int, int, int, int]:
|
||||||
|
"""
|
||||||
|
Expand a bounding box by a percentage margin, clamped to image bounds.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
box: Original bounding box
|
||||||
|
margin: Expansion factor (0.2 = 20% on each side)
|
||||||
|
image_size: (width, height) of the image
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(x1, y1, x2, y2) of expanded box
|
||||||
|
"""
|
||||||
|
w_img, h_img = image_size
|
||||||
|
bw = box.x2 - box.x1
|
||||||
|
bh = box.y2 - box.y1
|
||||||
|
|
||||||
|
x1 = max(0, int(box.x1 - bw * margin))
|
||||||
|
y1 = max(0, int(box.y1 - bh * margin))
|
||||||
|
x2 = min(w_img, int(box.x2 + bw * margin))
|
||||||
|
y2 = min(h_img, int(box.y2 + bh * margin))
|
||||||
|
|
||||||
|
return x1, y1, x2, y2
|
||||||
|
|
||||||
|
|
||||||
|
def point_in_box(point: Tuple[int, int], box: BoundingBox) -> bool:
|
||||||
|
"""Check if a point (x, y) is inside a bounding box."""
|
||||||
|
x, y = point
|
||||||
|
return box.x1 <= x <= box.x2 and box.y1 <= y <= box.y2
|
||||||
|
|
||||||
|
|
||||||
|
def box_center(box: BoundingBox) -> Tuple[int, int]:
|
||||||
|
"""Get center point of a bounding box."""
|
||||||
|
cx = (box.x1 + box.x2) // 2
|
||||||
|
cy = (box.y1 + box.y2) // 2
|
||||||
|
return cx, cy
|
||||||
|
|
||||||
|
|
||||||
|
def point_in_polygon(point: Tuple[int, int], polygon: List[Tuple[int, int]]) -> bool:
|
||||||
|
"""
|
||||||
|
Ray casting algorithm to check if point is inside polygon.
|
||||||
|
Useful for zone filtering (ROI).
|
||||||
|
"""
|
||||||
|
x, y = point
|
||||||
|
n = len(polygon)
|
||||||
|
inside = False
|
||||||
|
|
||||||
|
j = n - 1
|
||||||
|
for i in range(n):
|
||||||
|
xi, yi = polygon[i]
|
||||||
|
xj, yj = polygon[j]
|
||||||
|
|
||||||
|
if ((yi > y) != (yj > y)) and (x < (xj - xi) * (y - yi) / (yj - yi) + xi):
|
||||||
|
inside = not inside
|
||||||
|
j = i
|
||||||
|
|
||||||
|
return inside
|
||||||
|
|
||||||
|
|
||||||
|
def box_in_zone(box: BoundingBox, zone: List[Tuple[int, int]]) -> bool:
|
||||||
|
"""Check if box center is inside a zone polygon."""
|
||||||
|
center = box_center(box)
|
||||||
|
return point_in_polygon(center, zone)
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
"""Non-Maximum Suppression algorithms for detection filtering."""
|
||||||
|
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from pipeline.types import Detection
|
||||||
|
|
||||||
|
|
||||||
|
def nms(detections: List[Detection], iou_threshold: float = 0.5) -> List[Detection]:
|
||||||
|
"""
|
||||||
|
Non-Maximum Suppression: Remove overlapping boxes, keeping the highest confidence ones.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
detections: List of Detection objects
|
||||||
|
iou_threshold: Boxes with IoU > threshold are suppressed
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Filtered list of detections
|
||||||
|
"""
|
||||||
|
if not detections:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Sort by confidence (highest first)
|
||||||
|
sorted_dets = sorted(detections, key=lambda d: d.score, reverse=True)
|
||||||
|
|
||||||
|
keep = []
|
||||||
|
while sorted_dets:
|
||||||
|
# Take the best one
|
||||||
|
best = sorted_dets.pop(0)
|
||||||
|
keep.append(best)
|
||||||
|
|
||||||
|
# Filter out boxes that overlap too much with 'best'
|
||||||
|
remaining = []
|
||||||
|
for det in sorted_dets:
|
||||||
|
if best.box.iou(det.box) < iou_threshold:
|
||||||
|
remaining.append(det)
|
||||||
|
sorted_dets = remaining
|
||||||
|
|
||||||
|
return keep
|
||||||
|
|
||||||
|
|
||||||
|
def nms_by_class(
|
||||||
|
detections: List[Detection], iou_threshold: float = 0.5
|
||||||
|
) -> List[Detection]:
|
||||||
|
"""
|
||||||
|
Run NMS separately for each class label.
|
||||||
|
This prevents boxes of different classes from suppressing each other.
|
||||||
|
"""
|
||||||
|
if not detections:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Group by label
|
||||||
|
by_label: dict[str, List[Detection]] = {}
|
||||||
|
for det in detections:
|
||||||
|
by_label.setdefault(det.label, []).append(det)
|
||||||
|
|
||||||
|
# Run NMS per class
|
||||||
|
result = []
|
||||||
|
for dets in by_label.values():
|
||||||
|
result.extend(nms(dets, iou_threshold))
|
||||||
|
|
||||||
|
return result
|
||||||
@@ -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,17 @@
|
|||||||
|
"""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.metrics import MetricsSink
|
||||||
|
from pipeline.sinks.screenshots import ScreenshotSink
|
||||||
|
from pipeline.sinks.video import VideoFileStore
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"FileAlertStore",
|
||||||
|
"MetricsSink",
|
||||||
|
"ScreenshotSink",
|
||||||
|
"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,64 @@
|
|||||||
|
"""Live window display sink for real-time visualization."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from pipeline.ops import draw_detections, draw_frame_overlay
|
||||||
|
from pipeline.protocols import Sink
|
||||||
|
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,102 @@
|
|||||||
|
"""Metrics sink - tracks pipeline performance stats."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from collections import defaultdict
|
||||||
|
from typing import Dict
|
||||||
|
|
||||||
|
from pipeline.protocols import Sink
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class MetricsSink(Sink):
|
||||||
|
"""
|
||||||
|
Tracks pipeline performance metrics:
|
||||||
|
- FPS (frames per second)
|
||||||
|
- Detection counts by label
|
||||||
|
- Alert count
|
||||||
|
- Processing latency
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, log_interval: int = 100):
|
||||||
|
self.log_interval = log_interval
|
||||||
|
|
||||||
|
# Timing
|
||||||
|
self.start_time = time.time()
|
||||||
|
self.frame_count = 0
|
||||||
|
self.last_log_time = self.start_time
|
||||||
|
|
||||||
|
# Detection stats
|
||||||
|
self.detection_counts: Dict[str, int] = defaultdict(int)
|
||||||
|
self.alert_count = 0
|
||||||
|
|
||||||
|
# Latency tracking
|
||||||
|
self.frame_timestamps: list[float] = []
|
||||||
|
|
||||||
|
logger.info("MetricsSink initialized. Logging every %d frames.", log_interval)
|
||||||
|
|
||||||
|
def handle(self, frame: Frame):
|
||||||
|
self.frame_count += 1
|
||||||
|
now = time.time()
|
||||||
|
|
||||||
|
# Track latency (time from capture to processing)
|
||||||
|
latency = now - frame.timestamp
|
||||||
|
self.frame_timestamps.append(latency)
|
||||||
|
# Keep only last 100 for rolling average
|
||||||
|
if len(self.frame_timestamps) > 100:
|
||||||
|
self.frame_timestamps.pop(0)
|
||||||
|
|
||||||
|
# Count detections by label
|
||||||
|
for det in frame.detections:
|
||||||
|
self.detection_counts[det.label] += 1
|
||||||
|
|
||||||
|
# Count alerts
|
||||||
|
if frame.ticket:
|
||||||
|
self.alert_count += 1
|
||||||
|
|
||||||
|
# Periodic logging
|
||||||
|
if self.frame_count % self.log_interval == 0:
|
||||||
|
self._log_stats(now)
|
||||||
|
|
||||||
|
def _log_stats(self, now: float):
|
||||||
|
elapsed = now - self.start_time
|
||||||
|
fps = self.frame_count / elapsed if elapsed > 0 else 0
|
||||||
|
avg_latency = (
|
||||||
|
sum(self.frame_timestamps) / len(self.frame_timestamps)
|
||||||
|
if self.frame_timestamps
|
||||||
|
else 0
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"[METRICS] Frames: %d | FPS: %.1f | Latency: %.1fms | Alerts: %d",
|
||||||
|
self.frame_count,
|
||||||
|
fps,
|
||||||
|
avg_latency * 1000,
|
||||||
|
self.alert_count,
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
"""Print final summary."""
|
||||||
|
elapsed = time.time() - self.start_time
|
||||||
|
fps = self.frame_count / elapsed if elapsed > 0 else 0
|
||||||
|
avg_latency = (
|
||||||
|
sum(self.frame_timestamps) / len(self.frame_timestamps)
|
||||||
|
if self.frame_timestamps
|
||||||
|
else 0
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("=" * 50)
|
||||||
|
logger.info("PIPELINE METRICS SUMMARY")
|
||||||
|
logger.info("=" * 50)
|
||||||
|
logger.info("Total Frames: %d", self.frame_count)
|
||||||
|
logger.info("Runtime: %.1f seconds", elapsed)
|
||||||
|
logger.info("Average FPS: %.1f", fps)
|
||||||
|
logger.info("Avg Latency: %.1f ms", avg_latency * 1000)
|
||||||
|
logger.info("Total Alerts: %d", self.alert_count)
|
||||||
|
logger.info("-" * 50)
|
||||||
|
logger.info("Detections by Label:")
|
||||||
|
for label, count in sorted(self.detection_counts.items()):
|
||||||
|
logger.info(" %-25s %d", label, count)
|
||||||
|
logger.info("=" * 50)
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
"""Screenshot sink - saves evidence images on alerts."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from pipeline.ops.drawing import draw_alert_banner, draw_detections
|
||||||
|
from pipeline.protocols import Sink
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ScreenshotSink(Sink):
|
||||||
|
"""
|
||||||
|
Saves a screenshot when a ticket (alert) fires.
|
||||||
|
Useful for evidence/audit trails without full video recording.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, output_dir: str = "screenshots", annotate: bool = True):
|
||||||
|
self.output_dir = Path(output_dir)
|
||||||
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
self.annotate = annotate
|
||||||
|
self.count = 0
|
||||||
|
|
||||||
|
logger.info("ScreenshotSink initialized. Saving to %s", output_dir)
|
||||||
|
|
||||||
|
def handle(self, frame: Frame):
|
||||||
|
if not frame.ticket:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.count += 1
|
||||||
|
ticket = frame.ticket
|
||||||
|
|
||||||
|
# Build filename with timestamp and confidence
|
||||||
|
filename = f"alert_{ticket.frame_id}_{ticket.confidence:.2f}.jpg"
|
||||||
|
filepath = self.output_dir / filename
|
||||||
|
|
||||||
|
# Prepare image
|
||||||
|
if self.annotate:
|
||||||
|
img = frame.data.copy()
|
||||||
|
draw_detections(img, frame.detections)
|
||||||
|
draw_alert_banner(img, f"SMOKING {ticket.confidence:.0%}")
|
||||||
|
else:
|
||||||
|
img = frame.data
|
||||||
|
|
||||||
|
# Save
|
||||||
|
cv2.imwrite(str(filepath), img)
|
||||||
|
logger.info("Saved screenshot: %s", filepath)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
logger.info("ScreenshotSink closed. Total screenshots: %d", self.count)
|
||||||
@@ -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.ops import draw_detections, draw_frame_overlay
|
||||||
|
from pipeline.protocols import Sink
|
||||||
|
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.")
|
||||||
@@ -1,133 +0,0 @@
|
|||||||
import logging
|
|
||||||
import threading
|
|
||||||
import time
|
|
||||||
from queue import Empty, Queue
|
|
||||||
from typing import Iterator
|
|
||||||
|
|
||||||
import cv2
|
|
||||||
|
|
||||||
from pipeline.protocols import StreamSource
|
|
||||||
from pipeline.types import Frame
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
class FileStreamSource(StreamSource):
|
|
||||||
def __init__(self, path: str):
|
|
||||||
self.path = path
|
|
||||||
self._fps = 30.0
|
|
||||||
|
|
||||||
# Probe for FPS if it's a single file
|
|
||||||
if "*" not in self.path:
|
|
||||||
try:
|
|
||||||
cap = cv2.VideoCapture(self.path)
|
|
||||||
if cap.isOpened():
|
|
||||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
|
||||||
cap.release()
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("Could not probe FPS from %s: %s", path, e)
|
|
||||||
|
|
||||||
@property
|
|
||||||
def fps(self) -> float:
|
|
||||||
return self._fps
|
|
||||||
|
|
||||||
def stream(self) -> Iterator[Frame]:
|
|
||||||
# Handle glob patterns if path is a directory
|
|
||||||
if "*" in self.path:
|
|
||||||
import glob
|
|
||||||
|
|
||||||
files = glob.glob(self.path)
|
|
||||||
# Todo: handle image sequence
|
|
||||||
return
|
|
||||||
|
|
||||||
logger.info("Opening video file: %s", self.path)
|
|
||||||
cap = cv2.VideoCapture(self.path)
|
|
||||||
frame_id = 0
|
|
||||||
while cap.isOpened():
|
|
||||||
ret, frame = cap.read()
|
|
||||||
if not ret:
|
|
||||||
logger.info("End of video file reached.")
|
|
||||||
break
|
|
||||||
|
|
||||||
yield Frame(timestamp=time.time(), data=frame, frame_id=frame_id)
|
|
||||||
frame_id += 1
|
|
||||||
cap.release()
|
|
||||||
|
|
||||||
def stop(self):
|
|
||||||
"""No-op for file source (stream ends naturally)."""
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class ThreadedRTSPSource(StreamSource):
|
|
||||||
"""
|
|
||||||
Reads from RTSP stream or webcam in a background thread to prevent buffer buildup
|
|
||||||
and allow main processing loop to run at its own pace without stalling the decoding.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, source: str | int, queue_size: int = 5):
|
|
||||||
self.source = source
|
|
||||||
self.queue: Queue[Frame] = Queue(maxsize=queue_size)
|
|
||||||
self._fps = 30.0
|
|
||||||
|
|
||||||
# Probe FPS (with timeout safety ideally, but simple here)
|
|
||||||
try:
|
|
||||||
cap = cv2.VideoCapture(self.source)
|
|
||||||
if cap.isOpened():
|
|
||||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
|
||||||
cap.release()
|
|
||||||
except Exception:
|
|
||||||
logger.warning("Could not probe FPS from source.")
|
|
||||||
|
|
||||||
self.stopped = False
|
|
||||||
self.thread = threading.Thread(target=self._update, daemon=True)
|
|
||||||
self._start()
|
|
||||||
logger.info("Initialized ThreadedRTSPSource for %s", source)
|
|
||||||
|
|
||||||
@property
|
|
||||||
def fps(self) -> float:
|
|
||||||
return self._fps
|
|
||||||
|
|
||||||
def _start(self):
|
|
||||||
self.thread.start()
|
|
||||||
|
|
||||||
def _update(self):
|
|
||||||
logger.info("Reader thread started for %s", self.source)
|
|
||||||
cap = cv2.VideoCapture(self.source)
|
|
||||||
frame_id = 0
|
|
||||||
while not self.stopped:
|
|
||||||
if not cap.isOpened():
|
|
||||||
# Try reconnecting
|
|
||||||
logger.warning("Stream not open, retrying...")
|
|
||||||
time.sleep(1)
|
|
||||||
cap = cv2.VideoCapture(self.source)
|
|
||||||
continue
|
|
||||||
|
|
||||||
ret, frame = cap.read()
|
|
||||||
if not ret:
|
|
||||||
# Stream ended or failed
|
|
||||||
logger.warning("Failed to read frame from RTSP, retrying...")
|
|
||||||
time.sleep(0.1)
|
|
||||||
continue
|
|
||||||
|
|
||||||
# If queue is full, remove the oldest frame to keep latest
|
|
||||||
if self.queue.full():
|
|
||||||
try:
|
|
||||||
self.queue.get_nowait()
|
|
||||||
except Empty:
|
|
||||||
pass
|
|
||||||
|
|
||||||
self.queue.put(Frame(timestamp=time.time(), data=frame, frame_id=frame_id))
|
|
||||||
frame_id += 1
|
|
||||||
cap.release()
|
|
||||||
logger.info("RTSP Thread stopped.")
|
|
||||||
|
|
||||||
def stream(self) -> Iterator[Frame]:
|
|
||||||
while not self.stopped:
|
|
||||||
try:
|
|
||||||
yield self.queue.get(timeout=1)
|
|
||||||
except Empty:
|
|
||||||
continue
|
|
||||||
|
|
||||||
def stop(self):
|
|
||||||
self.stopped = True
|
|
||||||
self.thread.join()
|
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
"""Stream sources for the pipeline."""
|
||||||
|
|
||||||
|
from pipeline.sources.file import FileStreamSource
|
||||||
|
from pipeline.sources.stream import ThreadedStreamSource
|
||||||
|
|
||||||
|
# Alias for backwards compatibility
|
||||||
|
ThreadedRTSPSource = ThreadedStreamSource
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"FileStreamSource",
|
||||||
|
"ThreadedStreamSource",
|
||||||
|
"ThreadedRTSPSource", # Alias
|
||||||
|
]
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
"""File-based video source."""
|
||||||
|
|
||||||
|
import glob
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from typing import Iterator
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from pipeline.protocols import StreamSource
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class FileStreamSource(StreamSource):
|
||||||
|
"""Reads frames from a video file."""
|
||||||
|
|
||||||
|
def __init__(self, path: str):
|
||||||
|
self.path = path
|
||||||
|
self._fps = 30.0
|
||||||
|
|
||||||
|
# Probe for FPS if it's a single file
|
||||||
|
if "*" not in self.path:
|
||||||
|
try:
|
||||||
|
cap = cv2.VideoCapture(self.path)
|
||||||
|
if cap.isOpened():
|
||||||
|
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||||
|
cap.release()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Could not probe FPS from %s: %s", path, e)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def fps(self) -> float:
|
||||||
|
return self._fps
|
||||||
|
|
||||||
|
def stream(self) -> Iterator[Frame]:
|
||||||
|
# Handle glob patterns for image sequences
|
||||||
|
if "*" in self.path:
|
||||||
|
files = sorted(glob.glob(self.path))
|
||||||
|
logger.info("Loading %d images from pattern: %s", len(files), self.path)
|
||||||
|
for frame_id, file in enumerate(files):
|
||||||
|
img = cv2.imread(file)
|
||||||
|
if img is not None:
|
||||||
|
yield Frame(timestamp=time.time(), data=img, frame_id=frame_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info("Opening video file: %s", self.path)
|
||||||
|
cap = cv2.VideoCapture(self.path)
|
||||||
|
frame_id = 0
|
||||||
|
while cap.isOpened():
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if not ret:
|
||||||
|
logger.info("End of video file reached.")
|
||||||
|
break
|
||||||
|
|
||||||
|
yield Frame(timestamp=time.time(), data=frame, frame_id=frame_id)
|
||||||
|
frame_id += 1
|
||||||
|
cap.release()
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
"""No-op for file source (stream ends naturally)."""
|
||||||
|
pass
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
"""Threaded source for RTSP streams and webcams."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from queue import Empty, Queue
|
||||||
|
from typing import Iterator
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from pipeline.protocols import StreamSource
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ThreadedStreamSource(StreamSource):
|
||||||
|
"""
|
||||||
|
Reads from RTSP stream or webcam in a background thread.
|
||||||
|
|
||||||
|
Prevents buffer buildup and allows main processing loop to run
|
||||||
|
at its own pace without stalling the decoding.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, source: str | int, queue_size: int = 5):
|
||||||
|
self.source = source
|
||||||
|
self.queue: Queue[Frame] = Queue(maxsize=queue_size)
|
||||||
|
self._fps = 30.0
|
||||||
|
|
||||||
|
# Probe FPS
|
||||||
|
try:
|
||||||
|
cap = cv2.VideoCapture(self.source)
|
||||||
|
if cap.isOpened():
|
||||||
|
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||||
|
cap.release()
|
||||||
|
except cv2.error: # pylint: disable=catching-non-exception
|
||||||
|
logger.warning("Could not probe FPS from source.")
|
||||||
|
|
||||||
|
self.stopped = False
|
||||||
|
self.thread = threading.Thread(target=self._reader_loop, daemon=True)
|
||||||
|
self.thread.start()
|
||||||
|
logger.info("Initialized ThreadedStreamSource for %s", source)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def fps(self) -> float:
|
||||||
|
return self._fps
|
||||||
|
|
||||||
|
def _reader_loop(self):
|
||||||
|
"""Background thread that continuously reads frames."""
|
||||||
|
logger.info("Reader thread started for %s", self.source)
|
||||||
|
cap = cv2.VideoCapture(self.source)
|
||||||
|
frame_id = 0
|
||||||
|
|
||||||
|
while not self.stopped:
|
||||||
|
if not cap.isOpened():
|
||||||
|
logger.warning("Stream not open, retrying...")
|
||||||
|
time.sleep(1)
|
||||||
|
cap = cv2.VideoCapture(self.source)
|
||||||
|
continue
|
||||||
|
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if not ret:
|
||||||
|
logger.warning("Failed to read frame, retrying...")
|
||||||
|
time.sleep(0.1)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Drop oldest if queue full (keep latest frames)
|
||||||
|
if self.queue.full():
|
||||||
|
try:
|
||||||
|
self.queue.get_nowait()
|
||||||
|
except Empty:
|
||||||
|
pass
|
||||||
|
|
||||||
|
self.queue.put(Frame(timestamp=time.time(), data=frame, frame_id=frame_id))
|
||||||
|
frame_id += 1
|
||||||
|
|
||||||
|
cap.release()
|
||||||
|
logger.info("Reader thread stopped.")
|
||||||
|
|
||||||
|
def stream(self) -> Iterator[Frame]:
|
||||||
|
"""Yield frames from the queue."""
|
||||||
|
while not self.stopped:
|
||||||
|
try:
|
||||||
|
yield self.queue.get(timeout=1)
|
||||||
|
except Empty:
|
||||||
|
continue
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
"""Stop the reader thread and release resources."""
|
||||||
|
self.stopped = True
|
||||||
|
if self.thread.is_alive():
|
||||||
|
self.thread.join(timeout=2.0)
|
||||||
@@ -58,62 +58,6 @@ class Detection:
|
|||||||
metadata: dict = field(default_factory=dict)
|
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
|
@dataclass
|
||||||
class Ticket:
|
class Ticket:
|
||||||
"""Represents a potential alert event."""
|
"""Represents a potential alert event."""
|
||||||
|
|||||||
Reference in new issue
Block a user