Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import torch
 import os
 import numpy as np
+import torchvision
 from torchvision.models.detection import FasterRCNN
 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
@@ -22,20 +23,6 @@ def load_model(model_path, backbone_name, num_classes):
     model.to(device)
     model.eval()
     return model
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    if backbone_name == "resnet50":
-        model = torch.load(model_path, map_location=device)
-    elif backbone_name == "mobilenet":
-        model = torch.load(model_path, map_location=device)
-    model.to(device)
-    model.eval()
-    return model
-    if backbone_name == "resnet50":
-        model = torch.load(model_path)
-    elif backbone_name == "mobilenet":
-        model = torch.load(model_path)
-    model.eval()
-    return model
 
 resnet_model = load_model('fasterrcnnResnet.pth', 'resnet50', num_classes=6)
 mobilenet_model = load_model('fasterrcnnMobilenet.pth', 'mobilenet', num_classes=6)
@@ -101,4 +88,4 @@ video_interface = gr.Interface(
     title="Video Inference"
 )
 
-gr.TabbedInterface([image_interface, video_interface], ["Image", "Video"]).launch()
+gr.TabbedInterface([image_interface, video_interface], ["Image", "Video"]).launch()