Fix: version dependency

app.py

@@ -30,16 +30,17 @@ def parse_detections(results, yolo_version):
             boxes.append((xmin, ymin, xmax, ymax))
             colors.append(COLORS[category])
             names.append(name)
-    else:
-        # For YOLOv8 or other versions
-        for box in results[0].boxes:
-            xmin, ymin, xmax, ymax = box.xyxy.tolist()
-            confidence = box.conf.item()
-            class_id = int(box.cls.item())
-            if confidence > 0.2:
-                boxes.append((xmin, ymin, xmax, ymax))
-                colors.append(COLORS[class_id % len(COLORS)])
-                names.append(results[0].names[class_id])
+    elif yolo_version == "yolov8":
+        # For YOLOv8
+        for result in results:
+            for box in result.boxes:
+                xmin, ymin, xmax, ymax = box.xyxy.tolist()
+                confidence = box.conf.item()
+                class_id = int(box.cls.item())
+                if confidence > 0.2:
+                    boxes.append((xmin, ymin, xmax, ymax))
+                    colors.append(COLORS[class_id % len(COLORS)])
+                    names.append(result.names[class_id])
 
     return boxes, colors, names
 
@@ -83,13 +84,15 @@ def process_image(image, yolo_versions=["yolov5"]):
         model = load_yolo_model(yolo_version)
         
         # Run YOLO detection
-        results = model([rgb_img])
-        boxes, colors, names = parse_detections(results.xyxy[0], yolo_version)
+        results = model([rgb_img])  # Ensure this is a list containing one image
+        
+        # Parse detections using updated function
+        boxes, colors, names = parse_detections(results[0], yolo_version)
         
-        detections_img = draw_detections(boxes, colors, names.copy(), rgb_img.copy())
+        detections_img = draw_detections(boxes, colors.copy(), names.copy(), rgb_img.copy())
 
         # Grad-CAM visualization
-        target_layers = [model.model.model[-1]]  # Use the last layer as target layer for Grad-CAM
+        target_layers = [model.model.model[-1]]  # Use last layer as target layer for Grad-CAM
         cam = EigenCAM(model=model.model.model[-1], target_layers=target_layers)
         grayscale_cam = cam(tensor)[0]