Merge branch 'kp/fix_task_chain' of https://github.com/openvinotoolki…

…t/training_extensions into kp/fix_task_chain

src/otx/algo/utils/xai_utils.py

@@ -225,7 +225,7 @@ def _get_image_data_name(
     subset = datamodule.subsets[subset_name]
     item = subset.dm_subset[img_id]
     img = item.media_as(Image)
-    img_data, _ = subset._get_img_data_and_shape(img)  # noqa: SLF001
+    img_data, _, _ = subset._get_img_data_and_shape(img)  # noqa: SLF001
     image_save_name = "".join([char if char.isalnum() else "_" for char in item.id])
     return img_data, image_save_name
 

src/otx/core/data/dataset/base.py

@@ -146,7 +146,7 @@ def _get_img_data_and_shape(
         self,
         img: Image,
         roi: dict[str, Any] | None = None,
-    ) -> tuple[np.ndarray, tuple[int, int]]:
+    ) -> tuple[np.ndarray, tuple[int, int], dict[str, Any] | None]:
         key = img.path if isinstance(img, ImageFromFile) else id(img)
         roi_meta = None
 

src/otx/core/data/dataset/segmentation.py

@@ -205,9 +205,9 @@ def _get_item_impl(self, index: int) -> SegDataEntity | None:
         roi = item.attributes.get("roi", None)
         img_data, img_shape, roi_meta = self._get_img_data_and_shape(img, roi)
         if item.annotations:
-            ori_shape = roi_meta["orig_image_shape"] if roi else img_shape
+            ori_shape = roi_meta["orig_image_shape"] if roi_meta else img_shape
             extracted_mask = _extract_class_mask(item=item, img_shape=ori_shape, ignore_index=self.ignore_index)
-            if roi:
+            if roi_meta:
                 extracted_mask = extracted_mask[roi_meta["y1"] : roi_meta["y2"], roi_meta["x1"] : roi_meta["x2"]]
 
             masks = tv_tensors.Mask(extracted_mask[None])

src/otx/core/data/dataset/tile.py

@@ -461,7 +461,7 @@ def _get_item_impl(self, index: int) -> TileInstSegDataEntity:  # type: ignore[o
         """
         item = self.dm_subset[index]
         img = item.media_as(Image)
-        img_data, img_shape = self._get_img_data_and_shape(img)
+        img_data, img_shape, _ = self._get_img_data_and_shape(img)
 
         gt_bboxes, gt_labels, gt_masks, gt_polygons = [], [], [], []