Allow reading pose model metadata from local folder

aeon/io/reader.py

@@ -313,9 +313,12 @@ def read(self, file: Path) -> pd.DataFrame:
         """Reads data from the Harp-binarized tracking file."""
         # Get config file from `file`, then bodyparts from config file.
         model_dir = Path(*Path(file.stem.replace("_", "/")).parent.parts[-4:])
-        config_file_dir = Path(self._model_root) / model_dir
+        config_file_dir = file.parent / model_dir
         if not config_file_dir.exists():
-            raise FileNotFoundError(f"Cannot find model dir {config_file_dir}")
+            config_file_dir = Path(self._model_root) / model_dir
+            if not config_file_dir.exists():
+                raise FileNotFoundError(f"Cannot find model dir {config_file_dir}")
+
         config_file = self.get_config_file(config_file_dir)
         identities = self.get_class_names(config_file)
         parts = self.get_bodyparts(config_file)
@@ -350,7 +353,7 @@ def read(self, file: Path) -> pd.DataFrame:
             parts = unique_parts
 
         # Set new columns, and reformat `data`.
-        data = self.class_int2str(data, config_file)
+        data = self.class_int2str(data, identities)
         n_parts = len(parts)
         part_data_list = [pd.DataFrame()] * n_parts
         new_columns = pd.Series(["identity", "identity_likelihood", "part", "x", "y", "part_likelihood"])
@@ -407,18 +410,12 @@ def get_bodyparts(config_file: Path) -> list[str]:
         return parts
 
     @staticmethod
-    def class_int2str(data: pd.DataFrame, config_file: Path) -> pd.DataFrame:
+    def class_int2str(data: pd.DataFrame, classes: list[str]) -> pd.DataFrame:
         """Converts a class integer in a tracking data dataframe to its associated string (subject id)."""
-        if config_file.stem == "confmap_config":  # SLEAP
-            with open(config_file) as f:
-                config = json.load(f)
-            try:
-                heads = config["model"]["heads"]
-                classes = util.find_nested_key(heads, "classes")
-            except KeyError as err:
-                raise KeyError(f"Cannot find classes in {config_file}.") from err
-            for i, subj in enumerate(classes):
-                data.loc[data["identity"] == i, "identity"] = subj
+        identity = data["identity"].astype("string")
+        for i, subj in enumerate(classes):
+            identity.loc[data[identity.name] == i] = subj
+        data[identity.name] = identity
         return data
 
     @classmethod

aeon/util.py

@@ -14,7 +14,7 @@ def find_nested_key(obj: dict | list, key: str) -> Any:
             found = find_nested_key(v, key)
             if found:
                 return found
-    else:
+    elif obj is not None:
         for item in obj:
             found = find_nested_key(item, key)
             if found: