Add: support for yolov8

requirements.txt

@@ -6,4 +6,5 @@ pillow
 opencv-python
 git+https://github.com/jacobgil/pytorch-grad-cam.git
 gradio
-ultralytics
+ultralytics
+pip install yolov8-explainer

yolov8.py

@@ -1,19 +1,18 @@
-import torch
+from yolov8_explainer import YOLOv8Explainer
+from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image
-import torchvision.transforms as transforms
-from pytorch_grad_cam import EigenCAM
-from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
-from ultralytics import YOLO
 
+# Set random colors for detection bounding boxes
 COLORS = np.random.uniform(0, 255, size=(80, 3))
+
 def parse_detections_yolov8(results):
     boxes, colors, names = [], [], []
     detections = results.boxes
     for box in detections:
         confidence = box.conf[0].item()
-        if confidence < 0.2:
+        if confidence < 0.2:  # Filter out low-confidence detections
             continue
         xmin, ymin, xmax, ymax = map(int, box.xyxy[0].tolist())
         category = int(box.cls[0].item())
@@ -23,7 +22,6 @@ def parse_detections_yolov8(results):
         names.append(name)
     return boxes, colors, names
 
-
 def draw_detections(boxes, colors, names, img):
     for box, color, name in zip(boxes, colors, names):
         xmin, ymin, xmax, ymax = box
@@ -33,49 +31,24 @@ def draw_detections(boxes, colors, names, img):
                     lineType=cv2.LINE_AA)
     return img
 
-
-def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
-    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
-    cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
-    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
     model.to('cpu')
     model.eval()
 
+    # Initialize YOLOv8 Explainer
+    explainer = YOLOv8Explainer(model)
+
     # Run YOLO detection
     results = model(image)
     boxes, colors, names = parse_detections_yolov8(results[0])
     detections_img = draw_detections(boxes, colors, names, image.copy())
 
-    # Prepare input tensor for Grad-CAM
-    img_float = np.float32(image) / 255
-    transform = transforms.ToTensor()
-    tensor = transform(img_float).unsqueeze(0)
-
-    # Grad-CAM visualization
-    target_layers = [model.model.model[-2]]  # Adjust the target layer if required
-    cam_image, renormalized_cam_image = generate_cam_image(model.model, target_layers, tensor, image, boxes)
+    # Generate Grad-CAM visualization
+    cam_image, renormalized_cam_image = explainer.visualize(image, results)
 
-    # Combine results
+    # Combine original, Grad-CAM, and renormalized Grad-CAM images
     final_image = np.hstack((image, cam_image, renormalized_cam_image))
-    caption = "Results using YOLOv8"
-    return Image.fromarray(final_image), caption
+    caption = "Results using YOLOv8 and YOLOv8Explainer"
+    return Image.fromarray(final_image), caption