Fix arrow format reader for multiclass ROI case (#4145)

Fix arrow format reader for multiclass roi case
openvinotoolkit · Dec 5, 2024 · d9a3cd8 · d9a3cd8
1 parent 1608d9b
commit d9a3cd8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/otx/core/data/dataset/classification.py b/src/otx/core/data/dataset/classification.py
@@ -39,7 +39,10 @@ def _get_item_impl(self, index: int) -> MulticlassClsDataEntity | None:
             labels_ids = [
                 label["label"]["_id"] for label in roi["labels"] if label["label"]["domain"] == "CLASSIFICATION"
             ]
-            label_anns = [self.label_info.label_names.index(label_id) for label_id in labels_ids]
+            if self.data_format == "arrow":
+                label_anns = [self.label_info.label_ids.index(label_id) for label_id in labels_ids]
+            else:
+                label_anns = [self.label_info.label_names.index(label_id) for label_id in labels_ids]
         else:
             # extract labels from annotations
             label_anns = [ann.label for ann in item.annotations if isinstance(ann, Label)]