From fa530a78f6bd3568b095d8904bbb73c1f266bda0 Mon Sep 17 00:00:00 2001 From: "Kim, Sungchul" Date: Fri, 16 Aug 2024 17:22:37 +0900 Subject: [PATCH] Fix visual prompting dataset for empty nonzero mask (#3854) * Fix pytorchcv version * Fix for empty nonzero points --- src/otx/core/data/dataset/visual_prompting.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/otx/core/data/dataset/visual_prompting.py b/src/otx/core/data/dataset/visual_prompting.py index 74389d4fbf2..0047e9350fe 100644 --- 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