Fix masks merging

openvinotoolkit · Sep 19, 2024 · d778af8 · d778af8
1 parent 9535dee
commit d778af8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model_api/python/model_api/models/visual_prompting.py b/model_api/python/model_api/models/visual_prompting.py
@@ -286,7 +286,7 @@ def learn(
                     masks = _polygon_to_mask(inputs_decoder["polygon"], *original_shape)
                 else:
                     raise RuntimeError("Unsupported type of prompt")
-                ref_mask[masks] += 1
+                ref_mask = np.where(masks, 1, ref_mask)
             ref_mask = np.clip(ref_mask, 0, 1)
 
             ref_feat: np.ndarray | None = None