21

.gitattributes

@@ -1,2 +1,3 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
 *.jsonl filter=lfs diff=lfs merge=lfs -text

.vscode/settings.json

@@ -4,5 +4,6 @@
     "**/*.rpa": true,
     "**/*.rpymc": true,
     "**/cache/": true
-  }
+  },
+  "pylint.args": ["--generated-members", "cv2.*"]
 }

main.py

@@ -4,6 +4,7 @@ import sys
 
 from pipeline.core import SmokingDetectionPipeline
 from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
+from pipeline.outputs import VideoFileStore
 from pipeline.sources import FileStreamSource, ThreadedRTSPSource
 from pipeline.types import Ticket
 
@@ -35,7 +36,7 @@ def main():
     parser.add_argument(
         "--cig-model",
         type=str,
-        default="yolov8n.pt",
+        default="models/cig-det-best.pt",
         help="Path to Cigarette YOLO model",
     )
     parser.add_argument(
@@ -44,6 +45,12 @@ def main():
         default="models/cig-cls-best.pt",
         help="Path to Cigarette Classifier model",
     )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output.mp4",
+        help="Path to output video file (will also generate .jsonl)",
+    )
     parser.add_argument("--debug", action="store_true", help="Increase verbosity")
 
     args = parser.parse_args()
@@ -57,18 +64,22 @@ def main():
     )
     if args.debug:
         logging.getLogger().setLevel(logging.DEBUG)
-    
+
     # Initialize implementation-specific classes
     person_model = YoloPersonDetector(args.person_model)
     cigarette_model = YoloCigaretteDetector(args.cig_model)
     verifier = YoloClassifier(args.verify_model)
 
+    # Initialize Output Sink
+    video_sink = VideoFileStore(args.output)
+
     pipeline = SmokingDetectionPipeline(
         person_model=person_model,
         cigarette_model=cigarette_model,
         verifier_model=verifier,
         vlm_model=None,  # Optional
         store=ConsoleStore(),
+        sinks=[video_sink],
     )
 
     if args.rtsp:
@@ -82,9 +93,12 @@ def main():
     try:
         pipeline.process_stream(source)
     except KeyboardInterrupt:
-        if args.rtsp:
+        print("Stopped by user.")
+    finally:
+        if args.rtsp and hasattr(source, "stop"):
             source.stop()
-        print("Stopped.")
+        if video_sink:
+            video_sink.close()
 
 
 if __name__ == "__main__":

pipeline/core.py

@@ -22,12 +22,14 @@ class SmokingDetectionPipeline:
         verifier_model: Verifier,
         vlm_model: Optional[Verifier] = None,
         store: Optional[AlertStore] = None,
+        sinks: Optional[list] = 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()
+        self.sinks = sinks or []
 
         # Buffer for temporal context (optional, for VLM or sequence logic)
         self.frame_buffer = deque(maxlen=30)
@@ -45,14 +47,22 @@ class SmokingDetectionPipeline:
                 logger.info(f"Processed {count} frames...")
             self.frame_buffer.append(frame)
             self._process_frame(frame)
+
+            # Update sinks (Visualizers, etc)
+            for sink in self.sinks:
+                sink.handle(frame)
+
         logger.info("Stream processing finished.")
 
     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
+        # optimization: Resize or skip frames if inference is slow
         people = self.person_model.predict(frame.data)
 
+        # Add to frame detections for visualization
+        frame.detections.extend(people)
+
         if not people:
             return  # FAST EXIT
 
@@ -78,6 +88,34 @@ class SmokingDetectionPipeline:
             if cig_candidates:
                 candidates.append((person, cig_candidates))
 
+                # Check detections (translate to global coordinates)
+                for cig in cig_candidates:
+                    # Clone or modify the cig detection to be global
+                    # Assuming cig.box is relevant to person_crop
+                    # Global X = person.box.x1 + local_x
+                    # Global Y = person.box.y1 + local_y
+
+                    # Update in place might affect other logic if it assumes local,
+                    # but logic below (Expansion) seems to rely on cig.box being used on raw_person_crop
+                    # which implies local coordinates.
+                    # So we should create a COPY for visualization.
+
+                    from pipeline.types import BoundingBox, Detection
+
+                    global_cig = Detection(
+                        box=BoundingBox(
+                            x1=person.box.x1 + cig.box.x1,
+                            y1=person.box.y1 + cig.box.y1,
+                            x2=person.box.x1 + cig.box.x2,
+                            y2=person.box.y1 + cig.box.y2,
+                            conf=cig.box.conf,
+                        ),
+                        score=cig.score,
+                        label="candidate_cig",
+                        metadata=cig.metadata,
+                    )
+                    frame.detections.append(global_cig)
+
         if not candidates:
             return  # EXIT
 
@@ -136,6 +174,24 @@ class SmokingDetectionPipeline:
 
                 if score > 0.8:  # High confidence threshold
                     confirmed_hits.append((person, cig, score))
+
+                    # Add verified detection to frame for visualization
+                    from pipeline.types import BoundingBox, Detection
+
+                    verified_det = Detection(
+                        box=BoundingBox(
+                            x1=person.box.x1 + cig.box.x1,
+                            y1=person.box.y1 + cig.box.y1,
+                            x2=person.box.x1 + cig.box.x2,
+                            y2=person.box.y1 + cig.box.y2,
+                            conf=score,
+                        ),
+                        score=score,
+                        label="verified_smoking",
+                        metadata={"verifier_score": score},
+                    )
+                    frame.detections.append(verified_det)
+
                 elif score > 0.4 and self.vlm_model:
                     # Ambiguous case: Send to VLM?
                     vlm_score = self.vlm_model.verify(

pipeline/models.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, List
 
-from ultralytics import YOLO
+from ultralytics.models import YOLO
 
 from pipeline.protocols import Model, Verifier
 from pipeline.types import BoundingBox, Detection
@@ -69,10 +69,6 @@ class YoloClassifier(Verifier):
 
     def verify(self, input_data: Any) -> float:
         # input_data: numpy array (H, W, C) - Usually a crop
-        # Determine strict classification
-        # Assuming model returns a list of Results
-        # For classification, we look at probs
-
         if input_data.size == 0:
             return 0.0
 
@@ -86,10 +82,16 @@ class YoloClassifier(Verifier):
         if probs is None:
             return 0.0
 
-        # Assuming binary classification [negative, positive] or specific class index
-        # We need to know which class index corresponds to "smoking/cigarette"
-        # If it's binary trained: 0: Fire/Smoke/Cig, 1: No? Or vice versa.
-        # Let's assume class 0 is the target for now, or use probs.top1conf if we trust the label.
-
-        # Taking the highest confidence score
-        return float(probs.top1conf)
+        # User specified Class 0 is "smoke"
+        # We return the probability of class 0 specifically.
+        # Ultralytics probs.data is a tensor of scores for each class.
+        try:
+            return float(probs.data[0])
+        except (AttributeError, IndexError):
+            # Fallback if structure is unexpected, though data[0] is standard for classification
+            logger.warning(
+                "Could not retrieve Class 0 probability. Defaulting to top1conf if class 0."
+            )
+            if hasattr(probs, "top1") and probs.top1 == 0:
+                return float(probs.top1conf)
+            return 0.0

pipeline/outputs.py

@@ -0,0 +1,117 @@
+import json
+import logging
+from pathlib import Path
+from typing import List, Optional
+
+import cv2
+import numpy as np
+
+from pipeline.types import Detection, Frame
+
+logger = logging.getLogger(__name__)
+
+
+class VideoFileStore:
+    """
+    Writes processed frames to a video file with visualization overlays.
+    Also optionally writes metadata sidecar files.
+    """
+
+    def __init__(self, output_path: str, fps: int = 30, fourcc: str = "mp4v"):
+        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
+
+        # Open metadata file
+        self.meta_file = open(self.metadata_path, "w")
+        logger.info(
+            f"Initialized VideoFileStore writing to {output_path} and {self.metadata_path}"
+        )
+
+    def handle(self, frame: Frame):
+        """Process a frame: Draw boxes and write to video/log."""
+        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):
+        if self.writer:
+            self.writer.release()
+        if self.meta_file:
+            self.meta_file.close()
+        logger.info("VideoFileStore closed.")