Update app.py

app.py

@@ -11,9 +11,11 @@ class_map = ClassMap(['raccoon','banana'])
 
 size = 384
 
-model_2 = models.torchvision.retinanet.model(
-    backbone=models.torchvision.retinanet.backbones.resnet18_fpn (pretrained=True),
-    num_classes=len(class_map),
+model_type = models.mmdet.retinanet
+
+model_2 = model_type.model(
+    backbone= model_type.backbones.swin_t_p4_w7_fpn_1x_coco (pretrained=True),
+    num_classes=len(class_map)
 )
 
 # load from model_repo:
@@ -21,7 +23,7 @@ model_2 = models.torchvision.retinanet.model(
 # from huggingface_hub import hf_hub_download
 # hf_hub_download(repo_id="Alesteba/deep_model_02", filename="retinanet_racoon.pth")
 
-state_dict = torch.load('./retinanet_racoon.pth', map_location=torch.device('cpu'))
+state_dict = torch.load('./mmdet_racoon.pth', map_location=torch.device('cpu'))
 
 model_2.load_state_dict(state_dict)
 
@@ -33,7 +35,7 @@ def predict(img):
     
     img = PIL.Image.fromarray(img, "RGB")
     
-    pred_dict_2  = models.torchvision.retinanet.fastai.end2end_detect(
+    pred_dict_2  = model_type.fastai.end2end_detect(
     
         img, 
         infer_tfms,