Fix visual prompting dataset for empty nonzero mask (#3854)

* Fix pytorchcv version * Fix for empty nonzero points
openvinotoolkit · Aug 16, 2024 · fa530a7 · fa530a7
1 parent ac61247
commit fa530a7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/otx/core/data/dataset/visual_prompting.py b/src/otx/core/data/dataset/visual_prompting.py
@@ -90,7 +90,7 @@ def _get_item_impl(self, index: int) -> VisualPromptingDataEntity | None:
             if isinstance(annotation, dmPolygon):
                 mask = tvMask(polygon_to_bitmap([annotation], *img_shape)[0])
                 mask_points = torch.nonzero(mask)
-                if len(mask_points[0]) == 0:
+                if len(mask_points) == 0:
                     # skip very small region
                     continue
 
@@ -225,7 +225,7 @@ def _get_item_impl(self, index: int) -> ZeroShotVisualPromptingDataEntity | None
                 # generate prompts from polygon
                 mask = tvMask(polygon_to_bitmap([annotation], *img_shape)[0])
                 mask_points = torch.nonzero(mask)
-                if len(mask_points[0]) == 0:
+                if len(mask_points) == 0:
                     # skip very small region
                     continue