Fix: Revert -5.

app.py

@@ -45,9 +45,13 @@ def load_yolo_model(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)  # Load yolov5 small model
+        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', 'yolov5s', pretrained=True)  # YOLOv8 can be accessed via yolov5 repo
+        model = torch.hub.load('ultralytics/yolov5:v7.0', 'yolov5', pretrained=True)  # YOLOv8 is part of the yolov5 repo starting from v7.0
+    elif version == "yolov10":
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5m', pretrained=True)  # Placeholder for YOLOv10 (use an appropriate version if available)
     else:
         raise ValueError(f"Unsupported YOLO version: {version}")
     
@@ -55,6 +59,7 @@ 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)
@@ -73,12 +78,10 @@ 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)
-        
-        # In YOLOv5, the last convolutional layer in the backbone is typically the one to use
-        target_layers = [model.model.model[-2]]  # Assuming last layer of the model is a Conv layer
-        
+        target_layers = [model.model.model.model[-2]]  # Assumes last layer is used for Grad-CAM
+
         # Run YOLO detection
-        results = model(rgb_img)
+        results = model([rgb_img])
         boxes, colors, names = parse_detections(results)
         detections_img = draw_detections(boxes, colors, names, rgb_img.copy())
 
@@ -101,21 +104,20 @@ def process_image(image, yolo_versions=["yolov5"]):
 
     return result_images
 
-# Define the Gradio interface
 interface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.CheckboxGroup(
-            choices=["yolov3", "yolov5", "yolov8"],
+            choices=["yolov3", "yolov5", "yolov7", "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.",
-    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()