Debug: parsing detections

yolov8.py

@@ -52,16 +52,23 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
 def xai_yolov8n(image):
     model = YOLO('yolov8n.pt')  # Load YOLOv8n pre-trained weights
     model.eval()
-    model.cpu()
 
+    # Check if GPU is available and use it
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
     target_layers = [model.model.model[-2]]  # Grad-CAM target layer
     results = model([image])
-    boxes, colors, names = parse_detections(results)
+    if isinstance(results, list):
+        results = results[0]  # Extracting the first result (if list)
+
+    boxes, colors, names = parse_detections([results])  # Ensure results are passed as a list
     detections_img = draw_detections(boxes, colors, names, image.copy())
     img_float = np.float32(image) / 255
     transform = transforms.ToTensor()
-    tensor = transform(img_float).unsqueeze(0)
+    tensor = transform(img_float).unsqueeze(0).to(device)  # Ensure tensor is on the right device
     cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
     final_image = np.hstack((image, cam_image, renormalized_cam_image))
+
+    # Return final image and a caption
     caption = "Results using YOLOv8n"
     return Image.fromarray(final_image), caption