Fix: resolved model access methodology

app.py

@@ -12,7 +12,6 @@ from PIL import Image
 import gradio as gr
 from ultralytics import YOLO
 
-# Load a COCO-pretrained YOLOv3n model
 # Global Color Palette
 COLORS = np.random.uniform(0, 255, size=(80, 3))
 
@@ -47,7 +46,7 @@ def load_yolo_model(version="yolov5"):
     if version == "yolov5":
         model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
     elif version == "yolov8":
-        model = YOLO("yolov8n.pt")  # YOLOv8 is part of the yolov5 repo starting from v7.0
+        model = YOLO("yolov8n.pt")  # YOLOv8 is part of the ultralytics library
     elif version == "yolov10":
         model = torch.hub.load('ultralytics/yolov5', 'yolov5m', pretrained=True)  # Placeholder for YOLOv10 (use an appropriate version if available)
     else:
@@ -57,7 +56,6 @@ def load_yolo_model(version="yolov5"):
     model.cpu()
     return model
 
-# Main function for Grad-CAM visualization
 # Main function for Grad-CAM visualization
 def process_image(image, yolo_versions=["yolov5"]):
     image = np.array(image)
@@ -76,7 +74,7 @@ def process_image(image, yolo_versions=["yolov5"]):
     for yolo_version in yolo_versions:
         # 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
+        target_layers = [model.model[-1]]  # Fix: Access the correct last layer in YOLOv8
 
         # Run YOLO detection
         results = model([rgb_img])
@@ -102,20 +100,21 @@ 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",  "yolov8", "yolov10"],
+            choices=["yolov5", "yolov8", "yolov10"],
             value=["yolov5"],  # Set the default value (YOLOv5 checked by default)
             label="Select Model(s)",
         )
     ],
-    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.",
+    outputs=gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500),
+    title="Visualizing the key image features that drive decisions with our explainable AI tool.",
     description="XAI: Upload an image to visualize object detection of your models.."
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()