From d9a3cd8f06d1af1cade603dbe729c45efcbbbeb7 Mon Sep 17 00:00:00 2001 From: Vladislav Sovrasov Date: Thu, 5 Dec 2024 11:20:22 +0100 Subject: [PATCH] Fix arrow format reader for multiclass ROI case (#4145) Fix arrow format reader for multiclass roi case --- src/otx/core/data/dataset/classification.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/otx/core/data/dataset/classification.py b/src/otx/core/data/dataset/classification.py index f955c5bb8f..40b9c4a99f 100644 --- 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)]