add support for single session anipose

Author: darsnack

packages/cheese3d/cheese3d/cli.py

@@ -155,14 +155,17 @@ def calibrate(
     name: Annotated[str, typer.Argument(help="Name of project")] = ".",
     config_overrides: Annotated[Optional[List[str]], typer.Argument(
         help="Config overrides passed to Hydra (https://hydra.cc/docs/intro/)"
+    )] = None,
+    session: Annotated[Optional[str], typer.Option(
+        help="Name of a specific session subfolder to process"
     )] = None
 ):
     """Extract camera calibration for triangulation."""
     path = ctx.obj["path"]
     configs = ctx.obj["configs"]
     project = _build_project(path, name, configs, config_overrides)
     rich.print("Calibrating ...")
-    project.calibrate()
+    project.calibrate(session=session)
     rich.print("Calibration complete :white_check_mark:")
 
 @cli.command()
@@ -171,14 +174,17 @@ def track(
     name: Annotated[str, typer.Argument(help="Name of project")] = ".",
     config_overrides: Annotated[Optional[List[str]], typer.Argument(
         help="Config overrides passed to Hydra (https://hydra.cc/docs/intro/)"
+    )] = None,
+    session: Annotated[Optional[str], typer.Option(
+        help="Name of a specific session subfolder to process"
     )] = None
 ):
     """Track 2D keypoints using trained pose estimation model."""
     path = ctx.obj["path"]
     configs = ctx.obj["configs"]
     project = _build_project(path, name, configs, config_overrides)
     rich.print("Tracking 2D pose ...")
-    project.track()
+    project.track(session=session)
     rich.print("Pose estimation complete :white_check_mark:")
 
 @cli.command()
@@ -187,14 +193,17 @@ def triangulate(
     name: Annotated[str, typer.Argument(help="Name of project")] = ".",
     config_overrides: Annotated[Optional[List[str]], typer.Argument(
         help="Config overrides passed to Hydra (https://hydra.cc/docs/intro/)"
+    )] = None,
+    session: Annotated[Optional[str], typer.Option(
+        help="Name of a specific session subfolder to process"
     )] = None
 ):
     """Triangulate 2D pose estimation result to obtain 3D keypoints and Cheese3D features."""
     path = ctx.obj["path"]
     configs = ctx.obj["configs"]
     project = _build_project(path, name, configs, config_overrides)
     rich.print("Triangulating ...")
-    project.triangulate()
+    project.triangulate(session=session)
     rich.print("Triangulation complete :white_check_mark:")
 
 @cli.command()
@@ -203,6 +212,9 @@ def analyze(
     name: Annotated[str, typer.Argument(help="Name of project")] = ".",
     config_overrides: Annotated[Optional[List[str]], typer.Argument(
         help="Config overrides passed to Hydra (https://hydra.cc/docs/intro/)"
+    )] = None,
+    session: Annotated[Optional[str], typer.Option(
+        help="Name of a specific session subfolder to process"
     )] = None
 ):
     """
@@ -212,9 +224,9 @@ def analyze(
     path = ctx.obj["path"]
     configs = ctx.obj["configs"]
     project = _build_project(path, name, configs, config_overrides)
-    project.calibrate()
-    project.track()
-    project.triangulate()
+    project.calibrate(session=session)
+    project.track(session=session)
+    project.triangulate(session=session)
 
 @cli.command()
 def generate_videos(

packages/cheese3d/cheese3d/project.py

@@ -557,26 +557,54 @@ def _load_anipose_cfg(self):
 
         return load_config(str(self.triangulation_path / "config.toml"))
 
-    def calibrate(self):
-        from anipose.calibrate import calibrate_all
-        calibrate_all(self._load_anipose_cfg())
-
-    def track(self):
-        from anipose.pose_videos import pose_videos_all
-        pose_videos_all(self._load_anipose_cfg())
+    def _resolve_anipose_session(self, session: str) -> str:
+        session_path = self.triangulation_path / session
+        if not session_path.is_dir():
+            available = [p.name for p in self.triangulation_path.iterdir()
+                         if p.is_dir()] if self.triangulation_path.is_dir() else []
+            raise ValueError(
+                f"Session folder '{session}' not found in {self.triangulation_path}. "
+                f"Available sessions: {available}")
+        return str(session_path.resolve())
+
+    def calibrate(self, session: Optional[str] = None):
+        if session is not None:
+            from anipose.calibrate import process_session
+            config = self._load_anipose_cfg()
+            process_session(config, self._resolve_anipose_session(session))
+        else:
+            from anipose.calibrate import calibrate_all
+            calibrate_all(self._load_anipose_cfg())
+
+    def track(self, session: Optional[str] = None):
+        if session is not None:
+            from anipose.pose_videos import process_session
+            config = self._load_anipose_cfg()
+            process_session(config, self._resolve_anipose_session(session))
+        else:
+            from anipose.pose_videos import pose_videos_all
+            pose_videos_all(self._load_anipose_cfg())
 
-    def triangulate(self):
+    def triangulate(self, session: Optional[str] = None):
         # first triangulate points using Anipose
-        from anipose.triangulate import triangulate_all
-        triangulate_all(self._load_anipose_cfg())
+        if session is not None:
+            from anipose.triangulate import process_session
+            config = self._load_anipose_cfg()
+            process_session(config, self._resolve_anipose_session(session))
+            sessions_to_process = [self.triangulation_path / session]
+        else:
+            from anipose.triangulate import triangulate_all
+            triangulate_all(self._load_anipose_cfg())
+            sessions_to_process = [s for s in self.triangulation_path.iterdir()
+                                   if s.is_dir()]
         # now compute cheese3d features
         exclude_kps = set(kp.label for kp in _DEFAULT_KEYPOINTS) - set(kp.label for kp in self.keypoints)
         if len(exclude_kps) > 0:
             rprint("[bold red]Keypoint configuration does not match default Cheese3D keypoints. "
                    "Some Cheese3D features may not be computed![/bold red]")
-        for session in self.triangulation_path.iterdir():
-            if session.is_dir():
-                landmarks = read_3d_data(session)
+        for session_dir in sessions_to_process:
+            if session_dir.is_dir():
+                landmarks = read_3d_data(session_dir)
                 c3d_features = compute_anatomical_measurements(landmarks, exclude_kps)
                 # write features to csv
                 c3d_features_df = pd.DataFrame({
@@ -585,9 +613,9 @@ def triangulate(self):
                     for k, v in features.items()
                 })
                 if len(c3d_features_df) == 0:
-                    rprint(f"[bold red]No features constructed for {session.name}, skipping![/bold red]")
+                    rprint(f"[bold red]No features constructed for {session_dir.name}, skipping![/bold red]")
                     continue
-                csv_output = (session / "cheese3d")
+                csv_output = (session_dir / "cheese3d")
                 csv_output.mkdir(exist_ok=True)
                 csv_output = csv_output / "cheese3d_features.csv"
                 if not csv_output.exists():