refactor run_pipeline function for improved clarity and efficiency; add trigger class index handling and streamline detection logic

siwatsystem/pympta.py

@@ -27,10 +27,21 @@ def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict:
         model.to("cuda")
     else:
         logger.info(f"CUDA not available. Using CPU for model {node_config['modelId']}")
+
+    # Prepare trigger class indices for optimization
+    trigger_classes = node_config.get("triggerClasses", [])
+    trigger_class_indices = None
+    if trigger_classes and hasattr(model, "names"):
+        # Convert class names to indices for the model
+        trigger_class_indices = [i for i, name in model.names.items() 
+                                if name in trigger_classes]
+        logger.debug(f"Converted trigger classes to indices: {trigger_class_indices}")
+
     node = {
         "modelId": node_config["modelId"],
         "modelFile": node_config["modelFile"],
-        "triggerClasses": node_config.get("triggerClasses", []),
+        "triggerClasses": trigger_classes,
+        "triggerClassIndices": trigger_class_indices,
         "crop": node_config.get("crop", False),
         "minConfidence": node_config.get("minConfidence", None),
         "model": model,
@@ -158,130 +169,96 @@ def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
         logger.error(f"Error loading pipeline.json: {str(e)}", exc_info=True)
         return None
 
-def run_pipeline(frame, node: dict, return_bbox: bool = False, is_last_stage: bool = True):
+def run_pipeline(frame, node: dict, return_bbox: bool=False):
     """
-    Processes the frame with the given pipeline node. When return_bbox is True,
+    - For detection nodes (task != 'classify'):
-    the function returns a tuple (detection, bbox) where bbox is (x1,y1,x2,y2)
+        • runs `track(..., classes=triggerClassIndices)`
-    for drawing. Otherwise, returns only the detection.
+        • picks top box ≥ minConfidence
-    
+        • optionally crops & resizes → recurse into child
-    The is_last_stage parameter controls whether this node is considered the last
+        • else returns (det_dict, bbox)
-    in the pipeline chain. Only the last stage will return detection results.
+    - For classify nodes:
+        • runs `predict()`
+        • returns top (class,confidence) and no bbox
     """
     try:
-        # Check model type and use appropriate method
+        task = getattr(node["model"], "task", None)
-        model_task = getattr(node["model"], "task", None)
 
-        if model_task == "classify":
+        # ─── Classification stage ───────────────────────────────────
-            # Classification models need to use predict() instead of track()
+        if task == "classify":
-            logger.debug(f"Running classification model: {node.get('modelId')}")
+            # run the classifier and grab its top-1 directly via the Probs API
             results = node["model"].predict(frame, stream=False)
-            detection = None
+            # nothing returned?
-            best_box = None
+            if not results:
+                return (None, None) if return_bbox else None
 
-            # Process classification results
+            # take the first result's probs object
-            for r in results:
+            r     = results[0]
             probs = r.probs
-                if probs is not None and len(probs) > 0:
+            if probs is None:
-                    # Get the most confident class
+                return (None, None) if return_bbox else None
-                    class_id = int(probs.top1)
+
-                    conf = float(probs.top1conf)
+            # get the top-1 class index and its confidence
-                    detection = {
+            top1_idx  = int(probs.top1)
-                        "class": node["model"].names[class_id],
+            top1_conf = float(probs.top1conf)
-                        "confidence": conf,
+
-                        "id": None  # Classification doesn't have tracking IDs
+            det = {
+                "class": node["model"].names[top1_idx],
+                "confidence": top1_conf,
+                "id": None
             }
-                    logger.debug(f"Classification detection: {detection}")
+            return (det, None) if return_bbox else det
-                else:
-                    logger.debug(f"Empty classification results for model {node.get('modelId')}")
 
-            # Classification doesn't produce bounding boxes
-            bbox = None
 
-        else:
+        # ─── Detection stage ────────────────────────────────────────
-            # Detection/segmentation models use tracking
+        # only look for your triggerClasses
-            logger.debug(f"Running detection/tracking model: {node.get('modelId')}")
+        tk = node["triggerClassIndices"]
-            results = node["model"].track(frame, stream=False, persist=True)
+        res = node["model"].track(
-            detection = None
+            frame,
-            best_box = None
+            stream=False,
-            max_conf = -1
+            persist=True,
+            **({"classes": tk} if tk else {})
+        )[0]
 
-            # Log raw detection count
+        dets, boxes = [], []
-            detection_count = 0
+        for box in res.boxes:
-            for r in results:
+            conf = float(box.cpu().conf[0])
-                if hasattr(r.boxes, 'cpu') and len(r.boxes.cpu()) > 0:
+            cid  = int(box.cpu().cls[0])
-                    detection_count += len(r.boxes.cpu())
+            name = node["model"].names[cid]
+            if conf < node["minConfidence"]:
+                continue
+            xy = box.cpu().xyxy[0]
+            x1,y1,x2,y2 = map(int, xy)
+            dets.append({"class": name, "confidence": conf,
+                         "id": box.id.item() if hasattr(box, "id") else None})
+            boxes.append((x1, y1, x2, y2))
 
-            if detection_count == 0:
+        if not dets:
-                logger.debug(f"Empty detection results (no objects found) for model {node.get('modelId')}")
+            return (None, None) if return_bbox else None
-            else:
-                logger.debug(f"Detection model {node.get('modelId')} found {detection_count} objects")
 
-            for r in results:
+        # take highest‐confidence
-                for box in r.boxes:
+        best_idx = max(range(len(dets)), key=lambda i: dets[i]["confidence"])
-                    box_cpu = box.cpu()
+        best_det = dets[best_idx]
-                    conf = float(box_cpu.conf[0])
+        best_box = boxes[best_idx]
-                    if conf > max_conf and hasattr(box, "id") and box.id is not None:
-                        max_conf = conf
-                        detection = {
-                            "class": node["model"].names[int(box_cpu.cls[0])],
-                            "confidence": conf,
-                            "id": box.id.item()
-                        }
-                        best_box = box_cpu
 
-            if detection:
+        # ─── Branch (classification) ───────────────────────────────
-                logger.debug(f"Best detection: {detection}")
+        for br in node["branches"]:
-            else:
+            if (best_det["class"] in br["triggerClasses"]
-                logger.debug(f"No valid detection with tracking ID for model {node.get('modelId')}")
+                    and best_det["confidence"] >= br["minConfidence"]):
+                # crop if requested
+                sub = frame
+                if br["crop"]:
+                    x1,y1,x2,y2 = best_box
+                    sub = frame[y1:y2, x1:x2]
+                    sub = cv2.resize(sub, (224, 224))
 
-            bbox = None
+                det2, _ = run_pipeline(sub, br, return_bbox=True)
-            # Calculate bbox if best_box exists
+                if det2:
-            if detection and best_box is not None:
+                    # return classification result + original bbox
-                coords = best_box.xyxy[0]
+                    return (det2, best_box) if return_bbox else det2
-                x1, y1, x2, y2 = map(int, coords)
-                h, w = frame.shape[:2]
-                x1, y1 = max(0, x1), max(0, y1)
-                x2, y2 = min(w, x2), min(h, y2)
-                if x2 > x1 and y2 > y1:
-                    bbox = (x1, y1, x2, y2)
-                    logger.debug(f"Detection bounding box: {bbox}")
-                    if node.get("crop", False):
-                        frame = frame[y1:y2, x1:x2]
-                        logger.debug(f"Cropped frame to {frame.shape}")
 
-        # Check if we should process branches
+        # ─── No branch matched → return this detection ─────────────
-        if detection is not None:
+        return (best_det, best_box) if return_bbox else best_det
-            for branch in node["branches"]:
-                if detection["class"] in branch.get("triggerClasses", []):
-                    min_conf = branch.get("minConfidence")
-                    if min_conf is not None and detection["confidence"] < min_conf:
-                        logger.debug(f"Confidence {detection['confidence']} below threshold {min_conf} for branch {branch['modelId']}.")
-                        break
 
-                    # If we have branches, this is not the last stage
-                    branch_result = run_pipeline(frame, branch, return_bbox, is_last_stage=True)
-                    
-                    # This node is no longer the last stage, so its results shouldn't be returned
-                    is_last_stage = False
-                    
-                    if branch_result is not None:
-                        if return_bbox:
-                            return branch_result
-                        return branch_result
-                    break
-            
-            # Return this node's detection only if it's considered the last stage
-            if is_last_stage:
-                if return_bbox:
-                    return detection, bbox
-                return detection
-            
-        # No detection or not the last stage
-        if return_bbox:
-            return None, None
-        return None
     except Exception as e:
-        logger.error(f"Error running pipeline on node {node.get('modelId')}: {e}")
+        logging.error(f"Error in node {node.get('modelId')}: {e}")
-        if return_bbox:
+        return (None, None) if return_bbox else None
-            return None, None
-        return None