This commit is contained in:
lda committed 2026-01-07 00:01:55 +07:00
1 parent c1ee751099
commit 92b1aa2407
4 files changed
+43 -9

No files matched your search

+13 -2
View File
@@ -1,4 +1,5 @@
import argparse
import logging
import sys
from pipeline.core import SmokingDetectionPipeline
@@ -9,12 +10,19 @@ from pipeline.types import Ticket
class ConsoleStore:
def store(self, ticket: Ticket):
print(
logging.info(
f"[ALERT] {ticket.confidence:.2f} confirmed at {ticket.stage_confirmed} - Frame: {ticket.frame_id}"
)
def main():
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
handlers=[logging.StreamHandler(sys.stdout)],
)
parser = argparse.ArgumentParser(description="Smoking Detection Pipeline")
parser.add_argument("--source", type=str, help="Path to video file or RTSP stream")
parser.add_argument("--rtsp", action="store_true", help="Flag if source is RTSP")
@@ -36,6 +44,7 @@ def main():
default="models/cig-cls-best.pt",
help="Path to Cigarette Classifier model",
)
parser.add_argument("--debug", action="store_true", help="Increase verbosity")
args = parser.parse_args()
@@ -46,7 +55,9 @@ def main():
print(
f"Loading models... Person: {args.person_model}, Cig: {args.cig_model}, Verifier: {args.verify_model}"
)
if args.debug:
logging.getLogger().setLevel(logging.DEBUG)
# Initialize implementation-specific classes
person_model = YoloPersonDetector(args.person_model)
cigarette_model = YoloCigaretteDetector(args.cig_model)
+14 -7
View File
@@ -1,13 +1,15 @@
import logging
from collections import deque
from typing import Optional
from pipeline.protocols import AlertStore, Model, StreamSource, Verifier
from pipeline.types import Frame, Ticket
logger = logging.getLogger(__name__)
class ConsoleStore:
def store(self, ticket: Ticket):
print(
logger.info(
f"[STORE] Persisting Alert! Conf: {ticket.confidence:.2f} @ Stage: {ticket.stage_confirmed}"
)
@@ -29,15 +31,21 @@ class SmokingDetectionPipeline:
# Buffer for temporal context (optional, for VLM or sequence logic)
self.frame_buffer = deque(maxlen=30)
logger.info("Pipeline initialized.")
def process_stream(self, source: StreamSource):
"""Main loop processing the video stream."""
logger.info("Starting stream processing...")
# Batch processing strategy:
# Instead of strict batches, we process frame-by-frame but can look back
count = 0
for frame in source.stream():
count += 1
if count % 100 == 0:
logger.info(f"Processed {count} frames...")
self.frame_buffer.append(frame)
self._process_frame(frame)
logger.info("Stream processing finished.")
def _process_frame(self, frame: Frame):
# --- Step 1: Broad Phase (Person Detection) ---
@@ -48,6 +56,8 @@ class SmokingDetectionPipeline:
if not people:
return # FAST EXIT
logger.debug(f"Frame {frame.frame_id}: Found {len(people)} people.")
# --- Step 2: Narrow Phase (Cigarette Candidate) ---
# "Is the human holding something small?"
# We could optimize by only running on crops of the 'people'
@@ -71,6 +81,8 @@ class SmokingDetectionPipeline:
if not candidates:
return # EXIT
logger.debug(f"Frame {frame.frame_id}: Found {len(candidates)} candidates.")
# --- Step 3: Verification (Classifier) ---
# "Is that small thing actually a cigarette?"
confirmed_hits = []
@@ -90,11 +102,6 @@ class SmokingDetectionPipeline:
# Actually, we likely want to crop the cigarette from the person crop.
# Simplified: Just rely on the verifier taking the small patch.
cig_crop = cig.box.crop(
person.box.crop(frame.data)
) # Double crop if relative?
# If we assume cig.box is relative to person_crop:
cig_crop_img = cig.box.crop(frame.data) # This assumes global...
# To be robust, let's assume the Detector returns a Detection with a convenient 'crop' method or we handle coordinate translation.
# But for this MVP, let's assume we pass the crop derived from the box.
+6
View File
@@ -1,3 +1,4 @@
import logging
from typing import Any, List
from ultralytics import YOLO
@@ -5,9 +6,12 @@ from ultralytics import YOLO
from pipeline.protocols import Model, Verifier
from pipeline.types import BoundingBox, Detection
logger = logging.getLogger(__name__)
class YoloPersonDetector(Model):
def __init__(self, model_path: str = "yolo11n.pt"): # Using small model for speed
logger.info(f"Loading Person Detector: {model_path}")
self.model = YOLO(model_path) # Load model
def predict(self, input_data: Any) -> List[Detection]:
@@ -32,6 +36,7 @@ class YoloPersonDetector(Model):
class YoloCigaretteDetector(Model):
def __init__(self, model_path: str):
logger.info(f"Loading Cigarette Detector: {model_path}")
self.model = YOLO(model_path)
def predict(self, input_data: Any) -> List[Detection]:
@@ -59,6 +64,7 @@ class YoloCigaretteDetector(Model):
class YoloClassifier(Verifier):
def __init__(self, model_path: str):
logger.info(f"Loading Smoking Classifier: {model_path}")
self.model = YOLO(model_path)
def verify(self, input_data: Any) -> float:
+10
View File
@@ -1,3 +1,4 @@
import logging
import threading
import time
from queue import Empty, Queue
@@ -8,6 +9,8 @@ import cv2
from pipeline.protocols import StreamSource
from pipeline.types import Frame
logger = logging.getLogger(__name__)
class FileStreamSource(StreamSource):
def __init__(self, path: str):
@@ -22,11 +25,13 @@ class FileStreamSource(StreamSource):
# Todo: handle image sequence
return
logger.info(f"Opening video file: {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)
@@ -46,16 +51,19 @@ class ThreadedRTSPSource(StreamSource):
self.stopped = False
self.thread = threading.Thread(target=self._update, daemon=True)
self._start()
logger.info(f"Initialized ThreadedRTSPSource for {rtsp_url}")
def _start(self):
self.thread.start()
def _update(self):
logger.info(f"RTSP Thread started for {self.rtsp_url}")
cap = cv2.VideoCapture(self.rtsp_url)
frame_id = 0
while not self.stopped:
if not cap.isOpened():
# Try reconnecting
logger.warning("RTSP stream not open, retrying...")
time.sleep(1)
cap = cv2.VideoCapture(self.rtsp_url)
continue
@@ -63,6 +71,7 @@ class ThreadedRTSPSource(StreamSource):
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
@@ -76,6 +85,7 @@ class ThreadedRTSPSource(StreamSource):
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: