update webcam output

siwatsystem/pympta.py

@@ -87,36 +87,65 @@ def run_pipeline(frame, node: dict, return_bbox: bool = False):
     for drawing. Otherwise, returns only the detection.
     """
     try:
-        results = node["model"].track(frame, stream=False, persist=True)
-        detection = None
-        best_box = None
-        max_conf = -1
-
-        for r in results:
-            for box in r.boxes:
-                box_cpu = box.cpu()
-                conf = float(box_cpu.conf[0])
-                if conf > max_conf and hasattr(box, "id") and box.id is not None:
-                    max_conf = conf
+        # Check model type and use appropriate method
+        model_task = getattr(node["model"], "task", None)
+        
+        if model_task == "classify":
+            # Classification models need to use predict() instead of track()
+            logging.debug(f"Running classification model: {node.get('modelId')}")
+            results = node["model"].predict(frame, stream=False)
+            detection = None
+            best_box = None
+            
+            # Process classification results
+            for r in results:
+                probs = r.probs
+                if probs is not None and len(probs) > 0:
+                    # Get the most confident class
+                    class_id = int(probs.top1)
+                    conf = float(probs.top1conf)
                     detection = {
-                        "class": node["model"].names[int(box_cpu.cls[0])],
+                        "class": node["model"].names[class_id],
                         "confidence": conf,
-                        "id": box.id.item()
+                        "id": None  # Classification doesn't have tracking IDs
                     }
-                    best_box = box_cpu
+            
+            # Classification doesn't produce bounding boxes
+            bbox = None
+            
+        else:
+            # Detection/segmentation models use tracking
+            logging.debug(f"Running detection/tracking model: {node.get('modelId')}")
+            results = node["model"].track(frame, stream=False, persist=True)
+            detection = None
+            best_box = None
+            max_conf = -1
 
-        bbox = None
-        # Modified bounding box calculation: always compute bbox if best_box exists
-        if detection and best_box is not None:
-            coords = best_box.xyxy[0]
-            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)
-                if node.get("crop", False):
-                    frame = frame[y1:y2, x1:x2]
+            for r in results:
+                for box in r.boxes:
+                    box_cpu = box.cpu()
+                    conf = float(box_cpu.conf[0])
+                    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
+
+            bbox = None
+            # Calculate bbox if best_box exists
+            if detection and best_box is not None:
+                coords = best_box.xyxy[0]
+                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)
+                    if node.get("crop", False):
+                        frame = frame[y1:y2, x1:x2]
 
         if detection is not None:
             for branch in node["branches"]: