Properly load of the detection model

app.py

@@ -9,8 +9,6 @@ from model import get_model, predict, prepare_prediction, predict_class
 DET_CKPT = 'efficientDet_icevision.ckpt'
 CLASS_CKPT = 'class_ViT_taco_7_class.pth'
 
-
-
 st.subheader('Upload Custom Image')
 
 image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])

model.py

@@ -1,30 +1,28 @@
 from io import BytesIO
 from typing import Union
 from icevision import *
+from icevision.models.checkpoint import model_from_checkpoint
+from classifier import transform_image
+from icevision.models import ross
+
 import collections
 import PIL
 import torch
 import numpy as np
 import torchvision
 
-from classifier import transform_image
-
-import icevision.models.ross.efficientdet
-
-MODEL_TYPE = icevision.models.ross.efficientdet
+MODEL_TYPE = ross.efficientdet
 
 def get_model(checkpoint_path : str):
-    extra_args = {}
-    backbone = MODEL_TYPE.backbones.d0
-    # The efficientdet model requires an img_size parameter
-    extra_args['img_size'] = 512
-
-    model = MODEL_TYPE.model(backbone=backbone(pretrained=True),
-                             num_classes=2, 
-                             **extra_args)
-    
-    ckpt = get_checkpoint(checkpoint_path)
-    model.load_state_dict(ckpt)
+    checkpoint_and_model = model_from_checkpoint(
+                                checkpoint_path,
+                                model_name='ross.efficientdet',
+                                backbone_name='d0',
+                                img_size=512,
+                                classes=['Waste'],
+                                revise_keys=[(r'^model\.', '')])
+
+    model = checkpoint_and_model['model']
 
     return model
 
@@ -76,7 +74,7 @@ def prepare_prediction(pred_dict, threshold):
 
     return boxes, image
 
-def predict_class(model, image, bboxes):
+def predict_class(classifier, image, bboxes):
     preds = []
 
     for bbox in bboxes:
@@ -84,13 +82,12 @@ def predict_class(model, image, bboxes):
         bbox = np.array(bbox).astype(int)
         cropped_img = PIL.Image.fromarray(img).crop(bbox)
         cropped_img = np.array(cropped_img)
-        #cropped_img = torch.as_tensor(cropped_img, dtype=torch.float).unsqueeze(0)
 
         tran_image = transform_image(cropped_img, 224)
         tran_image = tran_image.transpose(2, 0, 1)
         tran_image = torch.as_tensor(tran_image, dtype=torch.float).unsqueeze(0)
         print(tran_image.shape)
-        y_preds = model(tran_image)
+        y_preds = classifier(tran_image)
         preds.append(y_preds.softmax(1).detach().numpy())
 
     preds = np.concatenate(preds).argmax(1)