Update: add support for other yolo variants

app.py

@@ -40,8 +40,27 @@ def draw_detections(boxes, colors, names, img):
                     lineType=cv2.LINE_AA)
     return img
 
+# Load the appropriate YOLO model based on the version
+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)
+    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)  # 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}")
+    
+    model.eval()  # Set to evaluation mode
+    model.cpu()
+    return model
+
 # Main function for Grad-CAM visualization
-def process_image(image):
+def process_image(image, yolo_versions=["yolov5"]):
     image = np.array(image)
     image = cv2.resize(image, (640, 640))
     rgb_img = image.copy()
@@ -51,41 +70,52 @@ def process_image(image):
     transform = transforms.ToTensor()
     tensor = transform(img_float).unsqueeze(0)
 
-    # Load YOLOv5 model
-    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-    model.eval()
-    model.cpu()
-    target_layers = [model.model.model.model[-2]]
+    # Initialize list to store result images
+    result_images = []
+
+    # Process each selected YOLO model
+    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
 
-    # Run YOLO detection
-    results = model([rgb_img])
-    boxes, colors, names = parse_detections(results)
-    detections_img = draw_detections(boxes, colors, names, rgb_img.copy())
+        # Run YOLO detection
+        results = model([rgb_img])
+        boxes, colors, names = parse_detections(results)
+        detections_img = draw_detections(boxes, colors, names, rgb_img.copy())
 
-    # Grad-CAM visualization
-    cam = EigenCAM(model, target_layers)
-    grayscale_cam = cam(tensor)[0, :, :]
-    cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
+        # Grad-CAM visualization
+        cam = EigenCAM(model, target_layers)
+        grayscale_cam = cam(tensor)[0, :, :]
+        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)
+        # 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)
 
-    # Concatenate images
-    final_image = np.hstack((rgb_img, cam_image, renormalized_cam_image))
+        # Concatenate images
+        final_image = np.hstack((rgb_img, cam_image, renormalized_cam_image))
+        result_images.append((yolo_version, Image.fromarray(final_image)))
 
-    return Image.fromarray(final_image)
+    return result_images
 
 # Gradio Interface
 interface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs=gr.Image(type="pil", label="Result"),
+    inputs=[
+        gr.Image(type="pil", label="Upload an Image"),
+        gr.CheckboxGroup(
+            choices=["yolov3", "yolov5", "yolov7", "yolov8", "yolov10"],
+            label="Select YOLO Models",
+            default=["yolov5"]
+        )
+    ],
+    outputs=gr.Gallery(label="Results", elem_id="gallery").style(grid=[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__":

requirements.txt

@@ -1,8 +1,12 @@
-torch
-torchvision
-torchaudio
-numpy
-pillow
-opencv-python
+torch==2.1.0
+torchvision==0.15.0
+torchaudio==2.1.0
+numpy==1.23.4
+pillow==9.3.0
+opencv-python==4.6.0.66
 git+https://github.com/jacobgil/pytorch-grad-cam.git
-gradio
+gradio==3.28.2
+git+https://github.com/ultralytics/yolov5.git  # For YOLOv5
+git+https://github.com/WongKinYiu/yolov7.git  # For YOLOv7
+git+https://github.com/ultralytics/ultralytics.git  # For YOLOv8
+git+https://github.com/saeedanwar/yolov10.git  # For YOLOv10