Update display.py

display.py

@@ -9,16 +9,16 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 from utils import YoloCAM, cells_to_bboxes, non_max_suppression
 from model import YOLOv3
 
-def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.5,show_cam: bool = False,  transparency: float = 0.5):
+def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.5,  transparency: float = 0.5):
     model = YOLOV3_PL() #YOLOv3(num_classes=20)
     model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
     # iou_thresh = 0.75
     # thresh = 0.75
     scaled_anchors  =  config.SCALED_ANCHORS
     
-    target_layer_list = model.layers[-2]
+   # target_layer_list = model.layers[-2]
 
-    cam = YoloCAM(model=model, target_layers = target_layer_list, use_cuda=False)
+   # cam = YoloCAM(model=model, target_layers = target_layer_list, use_cuda=False)
     transforms =  A.Compose(
     [
         A.LongestMaxSize(max_size=config.IMAGE_SIZE),
@@ -47,15 +47,16 @@ def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.5,sh
         bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
     )
     plot_img = draw_predictions(image, nms_boxes, class_labels=config.PASCAL_CLASSES)
-    if not show_cam:
-        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]
+    # 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]
 
+    return [plot_img]