Update ZSL outputs parser

openvinotoolkit · Aug 15, 2024 · 9d490f8 · 9d490f8
1 parent 3e9853d
commit 9d490f8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/otx/core/model/visual_prompting.py b/src/otx/core/model/visual_prompting.py
@@ -1041,9 +1041,9 @@ def _customize_outputs(  # type: ignore[override]
                 if len(predicted_mask.mask) == 0:
                     continue
                 _masks.append(np.stack(predicted_mask.mask, axis=0))
-                _used_points_scores = np.stack(predicted_mask.points, axis=0)
-                _prompts.append(_used_points_scores[:, :2])
-                _scores.append(_used_points_scores[:, 2])
+                _used_points_scores = np.stack(predicted_mask.scores, axis=0)
+                _prompts.append(np.stack(predicted_mask.points, axis=0))
+                _scores.append(_used_points_scores)
                 _labels.append(np.array([label] * len(_used_points_scores)))
 
             if len(_masks) == 0: