Update: support for yolo-explainer

app.py

@@ -10,6 +10,7 @@ from pytorch_grad_cam import EigenCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
 from PIL import Image
 import gradio as gr
+from YOLOv8_Explainer import yolov8_heatmap, display_images  # Import Explainer
 
 # Global Color Palette
 COLORS = np.random.uniform(0, 255, size=(80, 3))
@@ -42,8 +43,16 @@ def draw_detections(boxes, colors, names, img):
 
 # Load the appropriate YOLO model based on the version
 def load_yolo_model(version="yolov5"):
-    if version == "yolov5":
+    if version == "yolov3":
+        model = torch.hub.load('ultralytics/yolov3', 'yolov3', pretrained=True)
+    elif version == "yolov5":
         model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+    elif version == "yolov7":
+        model = torch.hub.load('WongKinYiu/yolov7', 'yolov7', pretrained=True)
+    elif version == "yolov8":
+        model = torch.hub.load('ultralytics/yolov5:v7.0', 'yolov5', pretrained=True)
+    elif version == "yolov10":
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5m', pretrained=True)
     else:
         raise ValueError(f"Unsupported YOLO version: {version}")
     
@@ -51,7 +60,8 @@ def load_yolo_model(version="yolov5"):
     model.cpu()
     return model
 
-def process_image(image, yolo_versions=["yolov5"]):
+# Main function for Grad-CAM visualization
+def process_image(image, yolo_versions=["yolov5"], use_explainer=False):
     image = np.array(image)
     image = cv2.resize(image, (640, 640))
     rgb_img = image.copy()
@@ -66,6 +76,23 @@ def process_image(image, yolo_versions=["yolov5"]):
 
     # Process each selected YOLO model
     for yolo_version in yolo_versions:
+        if use_explainer and yolo_version == "yolov8":
+            # Use YOLOv8 Explainer for EigenCAM heatmap
+            explainer_model = yolov8_heatmap(
+                weight="yolov8n.pt",
+                conf_threshold=0.4,
+                device="cpu",
+                method="EigenCAM",
+                layer=[10, 12, 14, 16, 18, -3],
+                backward_type="all",
+                ratio=0.02,
+                show_box=True,
+                renormalize=False,
+            )
+            imagelist = explainer_model(img_path=image)
+            display_images(imagelist)
+            continue  # Skip Grad-CAM for this case
+
         # Load the model based on YOLO version
         model = load_yolo_model(yolo_version)
         target_layers = [model.model.model.model[-2]]  # Assumes last layer is used for Grad-CAM
@@ -94,20 +121,22 @@ def process_image(image, yolo_versions=["yolov5"]):
 
     return result_images
 
+# Gradio Interface
 interface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.CheckboxGroup(
-            choices=["yolov5"],
-            value=["yolov5"],  # Set the default value (YOLOv5 checked by default)
+            choices=["yolov3", "yolov5", "yolov7", "yolov8", "yolov10"],
+            value=["yolov5"],  # Default to YOLOv5
             label="Select Model(s)",
-        )
+        ),
+        gr.Checkbox(label="Use YOLOv8 Explainer?", value=False)
     ],
-    outputs = gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500),
+    outputs=gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500),
     title="Visualising the key image features that drive decisions with our explainable AI tool.",
-    description="XAI: Upload an image to visualize object detection of your models.."
+    description="XAI: Upload an image to visualize object detection of your models."
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()