41

main.py

@@ -2,18 +2,11 @@ import argparse
 import sys
 
 from pipeline.core import SmokingDetectionPipeline
-from pipeline.models import YoloCigaretteDetector, YoloPersonDetector
-from pipeline.protocols import Verifier
+from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
 from pipeline.sources import FileStreamSource, ThreadedRTSPSource
 from pipeline.types import Ticket
 
 
-class MockVerifier(Verifier):
-    def verify(self, input_data):
-        # Stub for verification until we have that model
-        return 0.9
-
-
 class ConsoleStore:
     def store(self, ticket: Ticket):
         print(
@@ -28,7 +21,7 @@ def main():
     parser.add_argument(
         "--person-model",
         type=str,
-        default="yolov8n.pt",
+        default="models/human-det-best.pt",
         help="Path to Person YOLO model",
     )
     parser.add_argument(
@@ -37,6 +30,12 @@ def main():
         default="yolov8n.pt",
         help="Path to Cigarette YOLO model",
     )
+    parser.add_argument(
+        "--verify-model",
+        type=str,
+        default="models/cig-cls-best.pt",
+        help="Path to Cigarette Classifier model",
+    )
 
     args = parser.parse_args()
 
@@ -44,12 +43,14 @@ def main():
         print("Please provide --source")
         sys.exit(1)
 
-    print(f"Loading models... Person: {args.person_model}, Cig: {args.cig_model}")
+    print(
+        f"Loading models... Person: {args.person_model}, Cig: {args.cig_model}, Verifier: {args.verify_model}"
+    )
 
     # Initialize implementation-specific classes
     person_model = YoloPersonDetector(args.person_model)
     cigarette_model = YoloCigaretteDetector(args.cig_model)
-    verifier = MockVerifier()
+    verifier = YoloClassifier(args.verify_model)
 
     pipeline = SmokingDetectionPipeline(
         person_model=person_model,

pipeline/core.py

@@ -108,7 +108,19 @@ class SmokingDetectionPipeline:
 
                 # Re-cutting the crop:
                 raw_person_crop = person.box.crop(frame.data)
-                obj_crop = cig.box.crop(raw_person_crop)
+
+                # --- Expansion Logic ---
+                # Expand the box to give context
+                h, w = raw_person_crop.shape[:2]
+                pad_x = int((cig.box.x2 - cig.box.x1) * 0.5)  # Add 50% width padding
+                pad_y = int((cig.box.y2 - cig.box.y1) * 0.5)  # Add 50% height padding
+
+                x1 = max(0, cig.box.x1 - pad_x)
+                y1 = max(0, cig.box.y1 - pad_y)
+                x2 = min(w, cig.box.x2 + pad_x)
+                y2 = min(h, cig.box.y2 + pad_y)
+
+                obj_crop = raw_person_crop[y1:y2, x1:x2]
 
                 if obj_crop.size == 0:
                     continue

pipeline/models.py

@@ -2,7 +2,7 @@ from typing import Any, List
 
 from ultralytics import YOLO
 
-from pipeline.protocols import Model
+from pipeline.protocols import Model, Verifier
 from pipeline.types import BoundingBox, Detection
 
 
@@ -55,3 +55,35 @@ class YoloCigaretteDetector(Model):
                         )
                     )
         return detections
+
+
+class YoloClassifier(Verifier):
+    def __init__(self, model_path: str):
+        self.model = YOLO(model_path)
+
+    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
+
+        results = self.model(input_data, verbose=False)
+
+        # Check if we got any results
+        if not results:
+            return 0.0
+
+        probs = results[0].probs
+        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)