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