Update app.py

app.py

@@ -8,6 +8,20 @@ from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
 # Load Models
 def load_model(model_path, backbone_name, num_classes):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if backbone_name == "resnet50":
+        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+    elif backbone_name == "mobilenet":
+        model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(pretrained=False)
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+    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)