Update gradio_utils.py

gradio_utils.py

@@ -10,7 +10,7 @@ import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 from lightning_utils import YOLOv3Lightning
 from pytorch_grad_cam import GradCAM, EigenCAM
-from pytorch_grad_cam.utils.image import show_cam_on_image
+from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
 from pytorch_grad_cam.utils.model_targets import FasterRCNNBoxScoreTarget
 
 from utils import cells_to_bboxes, non_max_suppression
@@ -71,7 +71,9 @@ def infer_transform(IMAGE_SIZE=config.INFERENCE_IMAGE_SIZE):
             A.PadIfNeeded(
                 min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
             ),
-            A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
+            A.Normalize(mean=[0.45484068, 0.43406072, 0.40103856], 
+                        std=[0.23936155, 0.23471538, 0.23876129], 
+                        max_pixel_value=255,),
             ToTensorV2(),
         ]
   )
@@ -197,15 +199,24 @@ def upload_image_inference(img, transparency):
   bbox_coord = [b[2:] for b in nms_boxes_output[0]]
   targets = [FasterRCNNBoxScoreTarget(objs, bbox_coord)]
 
+  new_bboxes = [a[0] for a in annotations]
+
   cam = EigenCAM(model=model, 
                 target_layers=[model.model], 
                 reshape_transform=yolov3_reshape_transform)
 
   grayscale_cam = cam(input_tensor=img, targets=targets)
   grayscale_cam = grayscale_cam[0, :]
+
+  renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
+
+  for x1, y1, x2, y2 in new_bboxes:
+    renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())  
   
-  visualization = show_cam_on_image(img_copy/255, grayscale_cam, use_rgb=False, image_weight=transparency)
+  img_copy= np.float32(img_copy) / 255
+  renormalized_cam = scale_cam_image(renormalized_cam)
+  eigencam_image_renormalized = show_cam_on_image(img_copy, renormalized_cam, use_rgb=True)
 
   return([[img_copy, annotations],
-          [grayscale_cam, visualization]])
+          [grayscale_cam, eigencam_image_renormalized]])