Fix: parsing results

app.py

@@ -19,20 +19,29 @@ COLORS = np.random.uniform(0, 255, size=(80, 3))
 def parse_detections(results, yolo_version):
     if yolo_version == "yolov5":
         detections = results.pandas().xyxy[0].to_dict()
+        boxes, colors, names = [], [], []
+        for i in range(len(detections["xmin"])):
+            confidence = detections["confidence"][i]
+            if confidence < 0.2:
+                continue
+            xmin, ymin = int(detections["xmin"][i]), int(detections["ymin"][i])
+            xmax, ymax = int(detections["xmax"][i]), int(detections["ymax"][i])
+            name, category = detections["name"][i], int(detections["class"][i])
+            boxes.append((xmin, ymin, xmax, ymax))
+            colors.append(COLORS[category])
+            names.append(name)
     else:
-        detections = results[0].pandas().xyxy[0].to_dict()
+        boxes.append(results[0].boxes.xyxy)  # Bounding boxes in xyxy format (x1, y1, x2, y2)
+        confidences = results[0].boxes.conf  # Confidence scores
+        class_ids = results[0].boxes.cls  # Class IDs
+        names.append(results[0].names)  # Class names (from model)
+        colors = []
+        # Append predefined color based on category (class ID)
+        for class_id in class_ids:
+            # Map class ID to a color from the COLORS list (make sure you have enough colors)
+            color = COLORS[int(class_id) % len(COLORS)]  # Use modulo to avoid index error
+            colors.append(color)
 
-    boxes, colors, names = [], [], []
-    for i in range(len(detections["xmin"])):
-        confidence = detections["confidence"][i]
-        if confidence < 0.2:
-            continue
-        xmin, ymin = int(detections["xmin"][i]), int(detections["ymin"][i])
-        xmax, ymax = int(detections["xmax"][i]), int(detections["ymax"][i])
-        name, category = detections["name"][i], int(detections["class"][i])
-        boxes.append((xmin, ymin, xmax, ymax))
-        colors.append(COLORS[category])
-        names.append(name)
     return boxes, colors, names
 
 # Draw bounding boxes and labels