Update display.py

display.py

@@ -5,6 +5,8 @@ from utils import *
 import random
 from albumentations.pytorch import ToTensorV2
 from yolov3 import YOLOV3_PL
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from utils import YoloCAM, cells_to_bboxes, draw_predictions, non_max_suppression
 
 def inference(image: np.ndarray, iou_thresh: float = 0.75, thresh: float = 0.75,  transparency: float = 0.5):
     model = YOLOV3_PL()
@@ -40,7 +42,15 @@ def inference(image: np.ndarray, iou_thresh: float = 0.75, thresh: float = 0.75,
         bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
     )
     plot_img = draw_predictions(image, nms_boxes, class_labels=config.PASCAL_CLASSES)
-    return [plot_img]
+    if not show_cam:
+        return [plot_img]
+
+    grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]
+    img = cv2.resize(image, (416, 416))
+    img = np.float32(img) / 255
+    cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency)
+    return [plot_img, cam_image]
+