Cooked I Am

pipeline/models/__init__.py

@@ -0,0 +1,15 @@
+"""Model wrappers for the pipeline."""
+
+from pipeline.models.yolo import (
+    DEVICE,
+    YoloCigaretteDetector,
+    YoloClassifier,
+    YoloPersonDetector,
+)
+
+__all__ = [
+    "DEVICE",
+    "YoloPersonDetector",
+    "YoloCigaretteDetector",
+    "YoloClassifier",
+]

pipeline/models/yolo.py

@@ -0,0 +1,105 @@
+"""YOLO model wrappers for detection and classification."""
+
+import logging
+from typing import Any, List
+
+import torch
+from ultralytics.models import YOLO
+
+from pipeline.protocols import Model, Verifier
+from pipeline.types import BoundingBox, Detection
+
+logger = logging.getLogger(__name__)
+
+# Auto-detect device
+DEVICE = 0 if torch.cuda.is_available() else "cpu"
+if DEVICE == 0:
+    logger.info("GPU detected! Using CUDA device: %s", torch.cuda.get_device_name(0))
+else:
+    logger.info("No GPU detected. Running on CPU.")
+
+
+class YoloPersonDetector(Model):
+    """Detects people using YOLO."""
+
+    def __init__(self, model_path: str = "yolo11n.pt"):
+        logger.info("Loading Person Detector: %s", model_path)
+        self.model = YOLO(model_path)
+
+    def predict(self, input_data: Any) -> List[Detection]:
+        results = self.model(
+            input_data, classes=[0], verbose=False, device=DEVICE
+        )  # 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):
+    """Detects cigarettes using custom YOLO model."""
+
+    def __init__(self, model_path: str, confidence_threshold: float = 0.2):
+        logger.info("Loading Cigarette Detector: %s", model_path)
+        self.model = YOLO(model_path)
+        self.conf_threshold = confidence_threshold
+
+    def predict(self, input_data: Any) -> List[Detection]:
+        results = self.model(input_data, verbose=False, device=DEVICE)
+
+        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])
+                if conf > self.conf_threshold:
+                    detections.append(
+                        Detection(
+                            box=BoundingBox(x1, y1, x2, y2, conf),
+                            score=conf,
+                            label="cigarette_candidate",
+                        )
+                    )
+        return detections
+
+
+class YoloClassifier(Verifier):
+    """Classifies smoking using YOLO classification model."""
+
+    def __init__(self, model_path: str):
+        logger.info("Loading Smoking Classifier: %s", model_path)
+        self.model = YOLO(model_path)
+
+    def verify(self, input_data: Any) -> float:
+        if input_data.size == 0:
+            return 0.0
+
+        results = self.model(input_data, verbose=False, device=DEVICE)
+
+        if not results:
+            return 0.0
+
+        probs = results[0].probs
+        if probs is None:
+            return 0.0
+
+        # Class 0 is "smoke" - return its probability
+        try:
+            return float(probs.data[0])
+        except (AttributeError, IndexError):
+            logger.warning("Could not retrieve Class 0 probability. Checking top1.")
+            if hasattr(probs, "top1") and probs.top1 == 0:
+                return float(probs.top1conf)
+            return 0.0