Merge pull request #233 from KumarLabJax/pose-v6-segmentation-optional

make segmentation optional when loading v6+ pose files

Author: gbeane

src/jabs/feature_extraction/features.py

@@ -1,12 +1,13 @@
 from pathlib import Path
+from typing import cast
 
 import h5py
 import numpy as np
 import pandas as pd
 
 import jabs.project.track_labels
 from jabs.constants import COMPRESSION, COMPRESSION_OPTS_DEFAULT
-from jabs.pose_estimation import PoseEstimation, PoseHashException
+from jabs.pose_estimation import PoseEstimation, PoseEstimationV6, PoseHashException
 
 from .base_features import BaseFeatureGroup
 
@@ -113,7 +114,12 @@ def __init__(
         )
         self._cache_window = cache_window
         self._compute_social_features = pose_est.format_major_version >= 3
-        self._compute_segmentation_features = pose_est.format_major_version >= 6
+
+        self._compute_segmentation_features = (
+            pose_est.format_major_version >= 6
+            and cast(PoseEstimationV6, pose_est).has_segmentation
+        )
+
         distance_scale = (
             self._distance_scale_factor if self._distance_scale_factor is not None else 1.0
         )

src/jabs/pose_estimation/pose_est_v6.py

@@ -143,3 +143,8 @@ def get_segmentation_data_per_frame(self, frame_index, identity: int) -> np.ndar
     def format_major_version(self) -> int:
         """Returns the major version of the pose file format."""
         return 6
+
+    @property
+    def has_segmentation(self) -> bool:
+        """Returns True if segmentation data is available."""
+        return self._segmentation_dict["seg_data"] is not None