Add: support for yolov8

yolov8.py

@@ -35,23 +35,26 @@ def draw_detections(boxes, colors, names, img):
 
 
 def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
-    cam = EigenCAM(model, target_layers)
+    class YOLOWrapper(torch.nn.Module):
+        def __init__(self, model):
+            super(YOLOWrapper, self).__init__()
+            self.model = model
+
+        def forward(self, x):
+            return self.model.model.forward_once(x)  # Ensure correct layer is called
+
+    wrapped_model = YOLOWrapper(model)
+    cam = EigenCAM(wrapped_model, target_layers)
     grayscale_cam = cam(tensor)[0, :, :]
     img_float = np.float32(rgb_img) / 255
-
-    # Generate Grad-CAM
     cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
-
-    # Renormalize Grad-CAM inside bounding boxes
     renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
     for x1, y1, x2, y2 in boxes:
         renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())
     renormalized_cam = scale_cam_image(renormalized_cam)
     renormalized_cam_image = show_cam_on_image(img_float, renormalized_cam, use_rgb=True)
-
     return cam_image, renormalized_cam_image
 
-
 def xai_yolov8(image):
     # Load YOLOv8 model
     model = YOLO('yolov8n.pt')  # Load YOLOv8 nano model