fix kaggle support

openvinotoolkit · Jan 15, 2024 · e16ea18 · e16ea18
1 parent eb44355
commit e16ea18
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/src/datumaro/plugins/data_formats/kaggle/base.py b/src/datumaro/plugins/data_formats/kaggle/base.py
@@ -307,13 +307,13 @@ def __init__(
         for img_filename in os.listdir(path):
             if not img_filename.lower().endswith(tuple(IMAGE_EXTENSIONS)):
                 continue
-            item_id = os.path.splitext(img_filename)[0]
+            item_id = osp.splitext(img_filename)[0]
 
-            img_file = os.path.join(path, img_filename)
-            ann_file = os.path.join(ann_path, item_id + self.ann_extensions)
+            img_file = osp.join(path, img_filename)
+            ann_file = osp.join(ann_path, item_id + self.ann_extensions)
 
             annotations = (
-                self._parse_annotations(img_file, ann_file) if os.path.isfile(ann_file) else []
+                self._parse_annotations(img_file, ann_file) if osp.isfile(ann_file) else []
             )
 
             media = Image.from_file(path=img_file, size=self._size)
@@ -351,8 +351,11 @@ def _parse_annotations(self, img_file: str, ann_file: str):
             ymin = self._parse_field(bbox_elem, "ymin", float)
             ymax = self._parse_field(bbox_elem, "ymax", float)
 
-            self._label_cat.add(label_name)
-            label_id, _ = self._label_cat.find(label_name)
+            label_id, cat = self._label_cat.find(label_name)
+            if not cat:
+                self._label_cat.add(label_name)
+                label_id, _ = self._label_cat.find(label_name)
+
             annotations.append(
                 Bbox(id=obj_id, label=label_id, x=xmin, y=ymin, w=xmax - xmin, h=ymax - ymin)
             )
@@ -419,7 +422,10 @@ def _parse_annotations(self, img_file: str, ann_file: str):
             w *= image_width
             h *= image_height
 
-            self._label_cat.add(label_name)
+            label_id, cat = self._label_cat.find(label_name)
+            if not cat:
+                self._label_cat.add(label_name)
+                label_id, _ = self._label_cat.find(label_name)
             label_id, _ = self._label_cat.find(label_name)
 
             annotations.append(Bbox(id=obj_id, label=label_id, x=x, y=y, w=w, h=h))