Update app.py

app.py

@@ -17,6 +17,10 @@ import torch
 import torch.optim as optim
 import matplotlib
 import cv2
+from pytorch_grad_cam import EigenCAM
+from pytorch_grad_cam.utils.model_targets import FasterRCNNBoxScoreTarget
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
 
 # my files
 import utils
@@ -154,36 +158,43 @@ with gr.Blocks() as demo:
     colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
     colors_hex = {class_labels[i]:matplotlib.colors.rgb2hex(colors[i]) for i in range(0,len(class_labels))}
 
-    # app GUI
-    with gr.Row():
-        img_input  = gr.Image()
-        img_output = gr.AnnotatedImage().style(color_map = colors_hex)
-    section_btn = gr.Button("Identify Objects")
-
-    def yolo3_inference(input_img): # function for yolo inference
-
+    # consolidate the output from the model for gradcam to work
+    def yolov3_reshape_transform(x): 
+      activations = []
+      size = x[0].size()[2:4] # 13 * 13
+      for x_item in x:
+        x_permute = x_item.permute(0, 1, 4, 2, 3 ) # 1,3,25,13,13
+        x_permute = x_permute.reshape((x_permute.shape[0],
+                                    x_permute.shape[1]*x_permute.shape[2],
+                                    *x_permute.shape[3:])) # 1,75,13,13
+        activations.append(torch.nn.functional.interpolate(torch.abs(x_permute), size, mode='bilinear'))
+      activations = torch.cat(activations, axis=1) # 1,255,13,13
+      return(activations)
+        
+
+    # main function of the app
+    def yolo3_inference(input_img,gradcam=True,gradcam_opa=0.5): # function for yolo inference
+      
+      # load model
       yololit = LitYolo()
       inference_model = yololit.load_from_checkpoint("yolo3_model.ckpt")
+      
+      # bboxes, gradcam
       anchors  = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
       bboxes   = [[]]
       sections = [] # to return image and annotations
+      nms_boxes_output = []
+
+      # process the input image for inference/gradcam
+      input_img = cv2.resize(input_img, (416, 416))
+      input_img_copy = input_img.copy()
+      input_img = np.float32(input_img) / 255
+      transform = transforms.ToTensor()
+      input_img = transform(input_img).unsqueeze(0)
 
-      # image transformation
-      test_transforms = Al.Compose(
-        [
-            Al.LongestMaxSize(max_size=416),
-            Al.PadIfNeeded(
-                min_height=416, min_width=416, border_mode=cv2.BORDER_CONSTANT
-            ),
-            Al.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
-            ToTensorV2(),
-        ]
-      )
-      pr_input_img = test_transforms(image=input_img)
-      pr_input_img = pr_input_img['image'].unsqueeze(0)
       # infer the image
       inference_model.eval()
-      test_img_out   = inference_model(pr_input_img)
+      test_img_out   = inference_model(input_img)
 
       # process the outputs to create bounding boxes
       for i in range(3):
@@ -194,6 +205,7 @@ with gr.Blocks() as demo:
               bboxes[idx] += box
       # nms
       nms_boxes = utils.non_max_suppression(bboxes[0], iou_threshold=0.6, threshold=0.5, box_format="midpoint",)
+      nms_boxes_output.append(nms_boxes)
 
       # use gradio image annotations
       height, width = 416, 416
@@ -205,13 +217,37 @@ with gr.Blocks() as demo:
         lower_right_x = int(upper_left_x + (box[2] * width))
         lower_right_y = int(upper_left_y + (box[3] * height))
         sections.append(((upper_left_x,upper_left_y,lower_right_x,lower_right_y), class_labels[int(class_pred)]))
-      return (np.array(pr_input_img.squeeze(0).permute(1,2,0)),sections)
+      
+      # for gradcam
+      if gradcam:
+        objs = [b[1] for b in nms_boxes_output[0]]
+        bbox_coord = [b[2:] for b in nms_boxes_output[0]]
+        targets = [FasterRCNNBoxScoreTarget(objs, bbox_coord)]
+
+        target_layers = [inference_model.model]
+        cam = EigenCAM(inference_model, target_layers, use_cuda=False,reshape_transform=yolov3_reshape_transform)
+        grayscale_cam = cam(input_tensor = input_img, targets= targets)
+        grayscale_cam = grayscale_cam[0, :]
+        visualization = show_cam_on_image(input_img_copy/255, grayscale_cam, use_rgb=True, image_weight=gradcam_opa)
+
+        return (visualization,sections)
+      else:
+        return (np.array(input_img.squeeze(0).permute(1,2,0)),sections)
 
-    section_btn.click(yolo3_inference, inputs=[img_input], outputs=[img_output])
+    # app GUI
+    with gr.Row():
+        img_input  = gr.Image()
+        img_output = gr.AnnotatedImage().style(color_map = colors_hex)
+    with gr.Row():
+      gradcam_check = gr.Checkbox(label="Gradcam")
+      gradcam_opa = gr.Slider(0, 1, value = 0.5, label="Opacity of GradCAM")
+        
+    section_btn = gr.Button("Identify Objects")
+    section_btn.click(yolo3_inference, inputs=[img_input,gradcam_check,gradcam_opa], outputs=[img_output])
 
     gr.Markdown("## Some Examples")
     gr.Examples(examples=examples,
-                             inputs =img_input,
+                             inputs =[img_input,gradcam_check,gradcam_opa],
                              outputs=img_output,
                              fn=yolo3_inference, cache_examples=False)