93
This commit is contained in:
1 parent
c1ee751099
commit
92b1aa2407
4 files changed
+43
-9
No files matched your search
@@ -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
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user