Update display.py

display.py

@@ -9,8 +9,8 @@ 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.75, thresh: float = 0.75,show_cam: bool = False,  transparency: float = 0.5):
-    model = YOLOv3(num_classes=20)
+def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.5,show_cam: bool = False,  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