41

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)