67

pipeline/core.py

@@ -0,0 +1,139 @@
+from collections import deque
+from typing import Optional
+
+from pipeline.protocols import AlertStore, Model, StreamSource, Verifier
+from pipeline.types import Frame, Ticket
+
+
+class ConsoleStore:
+    def store(self, ticket: Ticket):
+        print(
+            f"[STORE] Persisting Alert! Conf: {ticket.confidence:.2f} @ Stage: {ticket.stage_confirmed}"
+        )
+
+
+class SmokingDetectionPipeline:
+    def __init__(
+        self,
+        person_model: Model,
+        cigarette_model: Model,
+        verifier_model: Verifier,
+        vlm_model: Optional[Verifier] = None,
+        store: Optional[AlertStore] = None,
+    ):
+        self.person_model = person_model
+        self.cigarette_model = cigarette_model
+        self.verifier_model = verifier_model
+        self.vlm_model = vlm_model
+        self.store = store or ConsoleStore()
+
+        # Buffer for temporal context (optional, for VLM or sequence logic)
+        self.frame_buffer = deque(maxlen=30)
+
+    def process_stream(self, source: StreamSource):
+        """Main loop processing the video stream."""
+
+        # Batch processing strategy:
+        # Instead of strict batches, we process frame-by-frame but can look back
+        for frame in source.stream():
+            self.frame_buffer.append(frame)
+            self._process_frame(frame)
+
+    def _process_frame(self, frame: Frame):
+        # --- Step 1: Broad Phase (Person Detection) ---
+        # "Is there a human here?"
+        # Optimization: Resize or skip frames if inference is slow
+        people = self.person_model.predict(frame.data)
+
+        if not people:
+            return  # FAST EXIT
+
+        # --- Step 2: Narrow Phase (Cigarette Candidate) ---
+        # "Is the human holding something small?"
+        # We could optimize by only running on crops of the 'people'
+        candidates = []
+        for person in people:
+            # ROI: Region of Interest (e.g., upper body or hands)
+            person_crop = person.box.crop(frame.data)
+
+            # If crop is too small, skip
+            if (
+                person_crop.size == 0
+                or person_crop.shape[0] < 10
+                or person_crop.shape[1] < 10
+            ):
+                continue
+
+            cig_candidates = self.cigarette_model.predict(person_crop)
+            if cig_candidates:
+                candidates.append((person, cig_candidates))
+
+        if not candidates:
+            return  # EXIT
+
+        # --- Step 3: Verification (Classifier) ---
+        # "Is that small thing actually a cigarette?"
+        confirmed_hits = []
+        for person, cigs in candidates:
+            for cig in cigs:
+                # Crop just the object.
+                # Note: box coordinates are relative to the person crop, or global?
+                # Ideally, detections should be unified.
+                # If cigarette_model returns boxes relative to person_crop, we need to adjust them.
+                # For simplicity in this mock, let's assume models return GLOBAL coords or handle their own relative logic wrapping.
+                # However, usually crop inference returns local coords.
+                # Let's assume for now the cig model knows how to return global.
+                # OR we pass the crop to verifier.
+
+                # IMPORTANT: If cigarette model ran on a crop, its output box is relative to that crop.
+                # We need to map it back or just use the crop directly?
+                # 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.
+
+                # FIX: If cig_candidates came from person_crop, their boxes are relative to (0,0) of person_crop.
+                # We should probably run verification on the crop defined by cig.box within person_crop.
+
+                # Let's grab the actual image patch for verification
+                # We can'teasily re-crop from global frame without translation.
+                # Strategy: cig detection on person_crop is easiest.
+
+                # Re-cutting the crop:
+                raw_person_crop = person.box.crop(frame.data)
+                obj_crop = cig.box.crop(raw_person_crop)
+
+                if obj_crop.size == 0:
+                    continue
+
+                score = self.verifier_model.verify(obj_crop)
+
+                if score > 0.8:  # High confidence threshold
+                    confirmed_hits.append((person, cig, score))
+                elif score > 0.4 and self.vlm_model:
+                    # Ambiguous case: Send to VLM?
+                    vlm_score = self.vlm_model.verify(
+                        obj_crop
+                    )  # Or pass seq of frames?
+                    if vlm_score > 0.8:
+                        confirmed_hits.append((person, cig, vlm_score))
+        # --- Step 4: Action ---
+        if confirmed_hits:
+            best_hit = max(confirmed_hits, key=lambda x: x[2])
+            ticket = Ticket(
+                frame_id=frame.frame_id,
+                confidence=best_hit[2],
+                evidence=[frame.data],  # Store frame reference
+                stage_confirmed="VLM"
+                if self.vlm_model
+                and best_hit[2] == 0.1  # Logic here is a bit shaky from mock
+                else "Verifier",
+            )
+            self.store.store(ticket)

pipeline/models.py

@@ -0,0 +1,57 @@
+from typing import Any, List
+
+from ultralytics import YOLO
+
+from pipeline.protocols import Model
+from pipeline.types import BoundingBox, Detection
+
+
+class YoloPersonDetector(Model):
+    def __init__(self, model_path: str = "yolo11n.pt"):  # Using small model for speed
+        self.model = YOLO(model_path)  # Load model
+
+    def predict(self, input_data: Any) -> List[Detection]:
+        # input_data: numpy array (H, W, C)
+        results = self.model(input_data, classes=[0], verbose=False)  # 0 is person
+
+        detections = []
+        for r in results:
+            boxes = r.boxes
+            for box in boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+                conf = float(box.conf[0])
+                detections.append(
+                    Detection(
+                        box=BoundingBox(x1, y1, x2, y2, conf),
+                        score=conf,
+                        label="person",
+                    )
+                )
+        return detections
+
+
+class YoloCigaretteDetector(Model):
+    def __init__(self, model_path: str):
+        self.model = YOLO(model_path)
+
+    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)
+
+        detections = []
+        for r in results:
+            boxes = r.boxes
+            for box in boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+                conf = float(box.conf[0])
+                # Maybe filter by conf here if needed
+                if conf > 0.2:
+                    detections.append(
+                        Detection(
+                            box=BoundingBox(x1, y1, x2, y2, conf),
+                            score=conf,
+                            label="cigarette_candidate",
+                        )
+                    )
+        return detections

pipeline/protocols.py

@@ -0,0 +1,27 @@
+from typing import Any, Iterator, List, Protocol
+
+from pipeline.types import Detection, Frame, Ticket
+
+
+class Model(Protocol):
+    def predict(self, input_data: Any) -> List[Detection]:
+        """Run inference and return detections."""
+        ...
+
+
+class Verifier(Protocol):
+    def verify(self, input_data: Any) -> float:
+        """Return a confidence score (0.0 to 1.0) for the specific target."""
+        ...
+
+
+class StreamSource(Protocol):
+    def stream(self) -> Iterator[Frame]:
+        """Yield frames from source."""
+        ...
+
+
+class AlertStore(Protocol):
+    def store(self, ticket: Ticket):
+        """Persist the ticket."""
+        ...

pipeline/sources.py

@@ -0,0 +1,89 @@
+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
+
+
+class FileStreamSource(StreamSource):
+    def __init__(self, path: str):
+        self.path = path
+
+    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
+
+        cap = cv2.VideoCapture(self.path)
+        frame_id = 0
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            yield Frame(timestamp=time.time(), data=frame, frame_id=frame_id)
+            frame_id += 1
+        cap.release()
+
+
+class ThreadedRTSPSource(StreamSource):
+    """
+    Reads from RTSP stream 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, rtsp_url: str, queue_size=5):
+        self.rtsp_url = rtsp_url
+        self.queue = Queue(maxsize=queue_size)
+        self.stopped = False
+        self.thread = threading.Thread(target=self._update, daemon=True)
+        self._start()
+
+    def _start(self):
+        self.thread.start()
+
+    def _update(self):
+        cap = cv2.VideoCapture(self.rtsp_url)
+        frame_id = 0
+        while not self.stopped:
+            if not cap.isOpened():
+                # Try reconnecting
+                time.sleep(1)
+                cap = cv2.VideoCapture(self.rtsp_url)
+                continue
+
+            ret, frame = cap.read()
+            if not ret:
+                # Stream ended or failed
+                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()
+
+    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()

pipeline/types.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass, field
+from typing import Any, List
+
+import numpy as np
+
+
+@dataclass
+class BoundingBox:
+    x1: int
+    y1: int
+    x2: int
+    y2: int
+    conf: float = 0.0
+
+    @property
+    def area(self):
+        return (self.x2 - self.x1) * (self.y2 - self.y1)
+
+    @property
+    def xyxy(self):
+        return [self.x1, self.y1, self.x2, self.y2]
+
+    def crop(self, image: Any) -> Any:
+        # Assuming image is a numpy array (H, W, C)
+        # Handle cases where box is partially out of bounds
+        h, w = image.shape[:2]
+        x1 = max(0, self.x1)
+        y1 = max(0, self.y1)
+        x2 = min(w, self.x2)
+        y2 = min(h, self.y2)
+        return image[y1:y2, x1:x2]
+
+
+@dataclass
+class Detection:
+    box: BoundingBox
+    score: float
+    label: str
+    metadata: dict = field(default_factory=dict)
+
+
+@dataclass
+class Frame:
+    """Represents a single frame in the video stream."""
+
+    timestamp: float
+    data: np.ndarray  # numpy array (OpenCV image)
+    frame_id: int
+    detections: List[Detection] = field(default_factory=list)
+
+
+@dataclass
+class Ticket:
+    """Represents a potential alert event."""
+
+    frame_id: int
+    confidence: float
+    evidence: List[Any]  # List of images or crops
+    stage_confirmed: str
+    metadata: dict = field(default_factory=dict)