Support point_labels 2, 3 (bbox left-top / right-bottom) in Prompt

wkentaro · May 5, 2024 · 609cbdd · 609cbdd
1 parent 4bab138
commit 609cbdd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/osam/types.py b/osam/types.py
@@ -52,8 +52,8 @@ def validate_point_labels(cls, point_labels, values):
             raise ValueError("point_labels must be 1-dimensional")
         if "points" in values and point_labels.shape[0] != values["points"].shape[0]:
             raise ValueError("point_labels must have the same number of rows as points")
-        if not set(np.unique(point_labels).tolist()).issubset({0, 1}):
-            raise ValueError("point_labels must contain only 0s and 1s")
+        if not set(np.unique(point_labels).tolist()).issubset({0, 1, 2, 3}):
+            raise ValueError("point_labels must contain only 0, 1, 2, or 3")
         return point_labels
 
     @pydantic.field_serializer("point_labels")