fix load checkpoint from path

app.py

@@ -10,7 +10,7 @@ from PIL import Image
 from streamlit_drawable_canvas import st_canvas
 from utils import utils
 
-PREDICTOR_MODEL, AUTOMASK_MODEL = utils.get_model('vit_b')
+PREDICTOR_MODEL, AUTOMASK_MODEL = utils.get_model('checkpoint/sam_vit_b_01ec64.pth')
 
 
 def process_box(predictor_model, show_mask, radius_width):

utils/utils.py

@@ -15,10 +15,11 @@ def get_color():
 
 
 @st.cache_resource
-def get_model(model):
+def get_model(model, checkpoint='checkpoint/sam_vit_b_01ec64.pth'):
     device         = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    build_sam      = sam_model_registry[model]
-    model          = build_sam(checkpoint=get_checkpoint_path(model)).to(device)
+    build_sam      = sam_model_registry['vit_b']
+    model          = build_sam(checkpoint=checkpoint)
+    model = model.to(device)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     predictor      = SamPredictor(model)