Add undefined label prompt check

Signed-off-by: heyufan1995 <[email protected]>

ci/unit_tests/test_vista3d.py

@@ -223,6 +223,26 @@
             "error": "Label prompt can only be a single object if provided with point prompts.",
         }
     ],
+    [
+        {
+            "bundle_root": "models/vista3d",
+            "input_dict": {"label_prompt": [16, 25, 26]},
+            "patch_size": [32, 32, 32],
+            "checkpointloader#_disabled_": True,  # do not load weights"
+            "initialize": ["$monai.utils.set_determinism(seed=123)"],
+            "error": "Undefined label prompt detected. Provide point prompts for zero-shot.",
+        }
+    ],
+    [
+        {
+            "bundle_root": "models/vista3d",
+            "input_dict": {"label_prompt": [136]},
+            "patch_size": [32, 32, 32],
+            "checkpointloader#_disabled_": True,  # do not load weights"
+            "initialize": ["$monai.utils.set_determinism(seed=123)"],
+            "error": "Undefined label prompt detected. Provide point prompts for zero-shot.",
+        }
+    ]
 ]
 
 

models/vista3d/scripts/evaluator.py

@@ -166,6 +166,10 @@ def check_prompts_format(self, label_prompt, points, point_labels):
                     raise ValueError("Label prompt must be a list of single scalar, [1,2,3,4,...,].")
                 if not np.all([(x < 255).item() for x in label_prompt]):
                     raise ValueError("Current bundle only supports label prompt smaller than 255.")
+                if points is None:
+                    supported_list = list(set([i+1 for i in range(132)]) - set([16,18,129,130,131]))
+                    if not np.all([x in supported_list for x in label_prompt]):
+                        raise ValueError("Undefined label prompt detected. Provide point prompts for zero-shot.")
             else:
                 raise ValueError("Label prompt must be a list, [1,2,3,4,...,].")
         # check points