Addin GradCAM

app.py

@@ -10,6 +10,8 @@ from PIL import Image
 from models.common import DetectMultiBackend
 from utils.augmentations import letterbox
 from utils.plots import Annotator, colors
+from pytorch_grad_cam import EigenCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
 from utils.torch_utils import select_device, smart_inference_mode
 from utils.general import check_img_size, Profile, non_max_suppression, scale_boxes
 
@@ -18,6 +20,7 @@ data = "data.yaml"
 # Load model
 device = select_device('cpu')
 model = DetectMultiBackend(weights, device=device, dnn=False, data=data, fp16=False)
+target_layers = [model.model[-2]]
 
 false_detection_data = glob(os.path.join("false_detection", '*.jpg'))
 false_detection_data = [x.replace('\\', '/') for x in false_detection_data]
@@ -61,6 +64,8 @@ def display_false_detection_data(false_detection_data, number_of_samples):
 
 def inference(input_img, conf_thres, iou_thres, is_false_detection_images=True, num_false_detection_images=10):
     im0 = input_img.copy()
+    im_resized = cv2.resize(im0, (640, 640))
+    rgb_img = im_resized.copy()
     stride, names, pt = model.stride, model.names, model.pt
     imgsz = check_img_size((640, 640), s=stride)  # check image size
 
@@ -107,7 +112,11 @@ def inference(input_img, conf_thres, iou_thres, is_false_detection_images=True,
     else:
         misclassified_images = None
 
-    return im0, misclassified_images
+    cam = EigenCAM(model, target_layers,task='od')
+    grayscale_cam = cam(rgb_img)[0, :, :]
+    cam_image = show_cam_on_image(im_resized, grayscale_cam, use_rgb=True)
+
+    return im0, cam_image, misclassified_images
 
 title = "YOLOv9 model to detect shirt/tshirt"
 description = "A simple Gradio interface to infer on YOLOv9 model and detect tshirt in image"
@@ -129,6 +138,7 @@ demo = gr.Interface(inference,
                               gr.Checkbox(label="Show False Detection"),
                               gr.Slider(5, 35, value=10, step=5, label="Number of False Detection")],
                     outputs= [gr.Image(width=640, height=640, label="Output"), 
+                              gr.Plot(label="GradCAM"),
                               gr.Plot(label="False Detection")],
                     title=title,
                     description=description,