SAMPredictor: adding CoralSCOP

Jordan-Pierce · Dec 11, 2024 · 15717ed · 15717ed
1 parent 44fa82f
commit 15717ed
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/coralnet_toolbox/SAM/QtDeployPredictor.py b/coralnet_toolbox/SAM/QtDeployPredictor.py
@@ -115,6 +115,7 @@ def setup_models_layout(self):
             "RepViT-SAM": "repvit.pt",
             "EdgeSAM": "edge_sam_3x.pt",
             "MobileSAM": "vit_t.pt",
+            "CoralSCOP": "vit_b_coralscop.pt",
             "SAM-Base": "vit_b.pt",
             "SAM-Large": "vit_l.pt",
             "SAM-Huge": "vit_h.pt"
@@ -295,6 +296,8 @@ def load_model(self):
                 model_type = "repvit"
             elif "edge_" in self.model_path.lower():
                 model_type = "edge_sam"
+            elif "_coralscop" in self.model_path.lower():
+                model_type = "vit_b_coralscop"
             elif "_t" in self.model_path.lower():
                 model_type = "vit_t"
             elif "_b" in self.model_path.lower():