Update display.py

display.py

@@ -4,9 +4,11 @@ import albumentations as A
 from utils import *
 import random
 from albumentations.pytorch import ToTensorV2
+from yolov3 import YOLOV3_PL
 
 def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.4, show_cam: bool = False, transparency: float = 0.5):
-
+    model = YOLOV3_PL()
+    model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
     transforms =  A.Compose(
     [
         A.LongestMaxSize(max_size=config.IMAGE_SIZE),