Add: support for yolov8

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import torchvision.transforms as transforms
 import gradio as gr
 from yolov5 import xai_yolov5
+from yolov8 import xai_yolov8
 
 def process_image(image, yolo_versions=["yolov5"]):
     image = np.array(image)
@@ -13,6 +14,8 @@ def process_image(image, yolo_versions=["yolov5"]):
     for yolo_version in yolo_versions:
         if yolo_version == "yolov5":
             result_images.append(xai_yolov5(image))
+        elif yolo_version == "yolov8":
+            result_images.append(xai_yolov8(image))
         else:
             result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
     return result_images
@@ -23,7 +26,7 @@ interface = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.CheckboxGroup(
-            choices=["yolov3", "yolov5", "yolov7", "yolov8", "yolov10"],
+            choices=["yolov5", "yolov8", "yolov10"],
             value=["yolov5"],  # Set default selection to YOLOv5
             label="Select Model(s)",
         )

yolov8.py

@@ -0,0 +1,78 @@
+import torch
+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
+
+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:
+            continue
+        xmin, ymin, xmax, ymax = map(int, box.xyxy[0].tolist())
+        category = int(box.cls[0].item())
+        name = results.names[category]
+        boxes.append((xmin, ymin, xmax, ymax))
+        colors.append(COLORS[category])
+        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
+        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
+        cv2.putText(img, name, (xmin, ymin - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2,
+                    lineType=cv2.LINE_AA)
+    return img
+
+
+def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
+    cam = EigenCAM(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
+    model.to('cpu')
+    model.eval()
+
+    # 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)
+
+    # Combine results
+    final_image = np.hstack((image, cam_image, renormalized_cam_image))
+    caption = "Results using YOLOv8"
+    return Image.fromarray(final_image), caption