Update app.py

app.py

@@ -10,11 +10,12 @@ import albumentations as A
 from albumentations.pytorch import ToTensorV2
 import config
 
-from utils import plot_single_image
+import utils 
 
 from model import YOLOv3
 import config
 from torchvision import transforms
+import torch.optim as optim
 
 scaled_anchors = (
             torch.tensor(config.ANCHORS)
@@ -22,7 +23,9 @@ scaled_anchors = (
         ).to('cpu')
 
 model = YOLOv3(num_classes=config.NUM_CLASSES)
-model.load_state_dict(torch.load("checkpoint.pth[1].tar", map_location=torch.device('cpu')), strict=False)
+optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY)
+
+utils.load_checkpoint("checkpoint.pth[1].tar", model, optimizer, config.LEARNING_RATE)
 
 test_transforms_exp = A.Compose(
     [
@@ -42,15 +45,16 @@ def inference(input_img,show_gradcam="yes", transparency = 0.5, target_layer_num
     input_img = transform(input_img)
     input_img = input_img
     input_img = input_img.unsqueeze(0)
-    out_fig = plot_single_image(model, input_img, 0.6, 0.5,scaled_anchors)
+    out_fig = utils.plot_single_image(model, input_img, 0.6, 0.5,scaled_anchors)
     return out_fig
 
 title = "TSAI S13 Assignment: YOLO V3 trained on PASCAL VOC Dataset"
-description = "A simple Gradio interface to infer on Custom ResNet model, and get GradCAM results. Please use images that belong to any of these classes - 'plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'."
-examples = [["cat.jpg","yes", 0.5, -1]            
+description = "A simple Gradio interface for object detection using YOLO V3 algorithm. Bounding boxes are shown around objects"
+examples = [["000002.jpg","yes", 0.5, -1],
+            ["000004.jpg","yes", 0.5, -1],
+            ["000006.jpg","yes", 0.5, -1],
+            ["000058.jpg","yes", 0.5, -1]
            ]
-
-#gr.Radio(["yes", "no"], label="Show Gradcam"),gr.Slider(0, 1, value = 0.5, label="If yes, Opacity of GradCAM"), gr.Slider(-2, -1, value = -2, step=1, label="If yes, Which Layer?")
 demo = gr.Interface(
     inference, 
     inputs = [gr.Image(shape=(416, 416), label="Input Image")],