fix native parser

luxonis · Jan 14, 2025 · beb4f2a · beb4f2a
1 parent 3a07ada
commit beb4f2a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 18 deletions.
diff --git a/luxonis_ml/data/parsers/base_parser.py b/luxonis_ml/data/parsers/base_parser.py
@@ -1,3 +1,4 @@
+import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
@@ -89,17 +90,20 @@ def _parse_split(self, **kwargs) -> List[Path]:
         @rtype: List[str]
         @return: List of added images.
         """
-        generator, skeletons, added_images = self.from_split(**kwargs)
-        self.dataset.add(self._add_task(generator))
-        if skeletons:
-            for skeleton in skeletons.values():
-                self.dataset.set_skeletons(
-                    skeleton.get("labels"),
-                    skeleton.get("edges"),
-                    self.dataset_type.value,
-                )
-
-        return added_images
+        old_cwd = os.getcwd()
+        try:
+            generator, skeletons, added_images = self.from_split(**kwargs)
+            self.dataset.add(self._add_task(generator))
+            if skeletons:
+                for skeleton in skeletons.values():
+                    self.dataset.set_skeletons(
+                        skeleton.get("labels"),
+                        skeleton.get("edges"),
+                        self.dataset_type.value,
+                    )
+            return added_images
+        finally:
+            os.chdir(old_cwd)
 
     def parse_split(
         self,

diff --git a/luxonis_ml/data/parsers/native_parser.py b/luxonis_ml/data/parsers/native_parser.py
@@ -1,4 +1,5 @@
 import json
+import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -65,14 +66,11 @@ def from_split(self, annotation_path: Path) -> ParserOutput:
             dictionary for keypoints and list of added images.
         """
 
+        data = json.loads(annotation_path.read_text())
+        os.chdir(annotation_path.parent)
+
         def generator() -> DatasetIterator:
-            for record in json.loads(annotation_path.read_text()):
-                if "file" not in record:
-                    raise ValueError(
-                        "The annotation record must contain 'file' key."
-                    )
-                record["file"] = annotation_path.parent / record["file"]
-                yield record
+            yield from data
 
         added_images = self._get_added_images(generator())