Add the two preliminary fine-tuned models

Author: constantinpape

micro_sam/util.py

@@ -25,13 +25,19 @@
 _MODEL_URLS = {
     "vit_h": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
     "vit_l": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
-    "vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+    "vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+    # preliminary finetuned models
+    "vit_h_lm": "https://owncloud.gwdg.de/index.php/s/CnxBvsdGPN0TD3A/download",
+    "vit_b_lm": "https://owncloud.gwdg.de/index.php/s/gGlR1LFsav0eQ2k/download",
 }
 _CHECKPOINT_FOLDER = os.environ.get("SAM_MODELS", os.path.expanduser("~/.sam_models"))
 _CHECKSUMS = {
     "vit_h": "a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
     "vit_l": "3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622",
-    "vit_b": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912"
+    "vit_b": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912",
+    # preliminary finetuned models
+    "vit_h_lm": "c30a580e6ccaff2f4f0fbaf9cad10cee615a915cdd8c7bc4cb50ea9bdba3fc09",
+    "vit_b_lm": "f2b8676f92a123f6f8ac998818118bd7269a559381ec60af4ac4be5c86024a1b",
 }
 
 
@@ -105,7 +111,13 @@ def get_sam_model(
     """
     checkpoint = _get_checkpoint(model_type, checkpoint_path)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    sam = sam_model_registry[model_type](checkpoint=checkpoint)
+
+    # Our custom model types have a suffix "_...". This suffix needs to be stripped
+    # before calling sam_model_registry.
+    model_type_ = model_type[:5]
+    assert model_type_ in ("vit_h", "vit_b", "vit_l")
+
+    sam = sam_model_registry[model_type_](checkpoint=checkpoint)
     sam.to(device=device)
     predictor = SamPredictor(sam)
     if return_sam: