Update app.py

app.py

@@ -8,7 +8,7 @@ from torch.utils.data import DataLoader
 import itertools
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-
+import cv2 
 
 import config as config
 from model import YOLOv3
@@ -17,7 +17,7 @@ from utils import get_loaders
 import utils
 
 new_state_dict = {}
-state_dict = torch.load('results/Yolov3_Lavanya.pth', map_location=torch.device('cpu'))
+state_dict = torch.load('Yolov3_Shashank.pth', map_location=torch.device('cpu'))
 for key, value in state_dict.items():
     new_key = key.replace('model.', '')
     new_state_dict[new_key] = value
@@ -48,7 +48,12 @@ classes = ("aeroplane",
     "tvmonitor")
 
 
-def inference(input_img=None, iou_threshold=0.6, conf_threshold=0.5):
+import grad_cam_func as gcf 
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+def inference(input_img=None, iou_threshold=0.6, conf_threshold=0.5, gc_trans=0.3):
     
     if input_img is not None:
         
@@ -133,26 +138,39 @@ def inference(input_img=None, iou_threshold=0.6, conf_threshold=0.5):
             
         ### GradCAM
         
-        # target_layer = [model.layers[-2]]
-        # target_categories = [box[0] for box in nms_boxes]
-        # targets = [ClassifierOutputTarget(category) for category in target_categories]
+        target_layer = [model.layers[-2]]
+        cam = gcf.BaseCAM(model, target_layer)
         
-        # cam = ex.BaseCAM(model, target_layer, target_categories)
-        # help_ = cam.compute_cam_per_layer(transform_img, targets, False)
+        AnG = ActivationsAndGradients(model, target_layer, None)
+        outputs = AnG(transform_img)
         
-        # print(help_)
-        # print(np.shape(cam))
-        
-        # cam = EigenCAM(model, [model.layers[-2]], use_cuda=False)
-        # grayscale_cam = cam(transform_img)[0, :, :]
-        # print('Hello')
-        # cam_image = show_cam_on_image(input_img, grayscale_cam, use_rgb=True)
+        bboxes = [[] for _ in range(1)]
+        for i in range(3):
+            batch_size, A, S, _, _ = outputs[i].shape
+            anchor = config.SCALED_ANCHORS[i]
+            boxes_scale_i = utils.cells_to_bboxes(
+                outputs[i], anchor, S=S, is_preds=True
+            )
+            for idx, (box) in enumerate(boxes_scale_i):
+                bboxes[idx] += box
         
-        # plt.imshow(cam_image)
-        # plt.show()
+        nms_boxes = utils.non_max_suppression(
+            bboxes[0], iou_threshold=0.5, threshold=0.4, box_format="midpoint",
+        )
 
+        target_categories = [box[0] for box in nms_boxes]
+        targets = [ClassifierOutputTarget(
+            category) for category in target_categories]
         
-        outputs_inference_gc = None
+        help_ = cam.compute_cam_per_layer(transform_img, targets, False)
+        
+        output_gc = cam.aggregate_multi_layers(help_)[0, :, :]
+        
+        img = cv2.resize(input_img, (416, 416))
+        img = np.float32(img) / 255
+        cam_image = show_cam_on_image(img, output_gc, use_rgb=True, image_weight=gc_trans)   
+             
+        outputs_inference_gc = cam_image
         
     else: 
         outputs_inference_bb = None 
@@ -161,14 +179,16 @@ def inference(input_img=None, iou_threshold=0.6, conf_threshold=0.5):
     return outputs_inference_bb, outputs_inference_gc
 
 
+
 title = "PASCAL VOC trained on Yolov3"
 description = "A simple Gradio interface to infer on Yolov3 model, and get GradCAM results"
-examples = [['examples/test_'+str(i)+'.jpg', 0.6, 0.5] for i in range(10)]
+examples = [['examples/test_'+str(i)+'.jpg', 0.6, 0.5, 0.3] for i in range(10)]
 
 demo = gr.Interface(inference,
                     inputs = [gr.Image(label="Input image"),
                                 gr.Slider(0, 1, value=0.6, label="IOU Threshold"),
                                 gr.Slider(0, 1, value=0.4, label="Threshold"),
+                                gr.Slider(0, 1, value=0.5, label="GradCAM Transparency"),
                               ],
                     outputs = [
                         gr.Image(label="Yolov3 Prediction"),