Fix arrow format reader for multiclass roi case

openvinotoolkit · Dec 4, 2024 · 57c47bf · 57c47bf
1 parent 5d6f8d3
commit 57c47bf
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)]