Improve rototranslate and get_body_model APIs

Make rototranslate t default to zero for pure-rotation use case.
Remove redundant DATA_ROOT/model_root defaulting from get_body_model
(now handled in BodyModel.__init__).

Author: isarandi

src/smplfitter/jax/__init__.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import functools
-import os
 
 from .bodymodel import BodyModel
 from .bodyfitter import BodyFitter
@@ -21,7 +20,4 @@ def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
 
 
 def get_body_model(model_name, gender, model_root=None):
-    if model_root is None:
-        DATA_ROOT = os.getenv('DATA_ROOT', default='.')
-        model_root = f'{DATA_ROOT}/body_models/{model_name}'
     return BodyModel(model_root=model_root, gender=gender, model_name=model_name)

src/smplfitter/jax/bodymodel.py

@@ -156,12 +156,23 @@ def __call__(
         )
 
     def rototranslate(
-        self, R, t, pose_rotvecs, shape_betas, trans, kid_factor=0, post_translate=True
+        self,
+        R,
+        t=None,
+        pose_rotvecs=None,
+        shape_betas=None,
+        trans=None,
+        kid_factor=0,
+        post_translate=True,
     ):
         """Rotate and translate the body in parametric form.
 
         See np.BodyModel.rototranslate for full documentation.
         """
+        if t is None:
+            t = jnp.zeros(3, dtype=R.dtype)
+        if pose_rotvecs is None or shape_betas is None or trans is None:
+            raise ValueError('pose_rotvecs, shape_betas, and trans are required.')
         current_rotmat = rotvec2mat(pose_rotvecs[:3])
         new_rotmat = R @ current_rotmat
         new_pose_rotvec = jnp.concatenate([mat2rotvec(new_rotmat), pose_rotvecs[3:]], axis=0)

src/smplfitter/nb/__init__.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import functools
-import os
 
 from .bodymodel import BodyModel
 from .bodyfitter import BodyFitter
@@ -21,7 +20,4 @@ def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
 
 
 def get_body_model(model_name, gender, model_root=None):
-    if model_root is None:
-        DATA_ROOT = os.getenv('DATA_ROOT', default='.')
-        model_root = f'{DATA_ROOT}/body_models/{model_name}'
     return BodyModel(model_root=model_root, gender=gender, model_name=model_name)

src/smplfitter/nb/bodymodel.py

@@ -223,7 +223,14 @@ def single(self, *args, return_vertices=True, **kwargs):
         return {k: np.squeeze(v, axis=0) for k, v in result.items()}
 
     def rototranslate(
-        self, R, t, pose_rotvecs, shape_betas, trans, kid_factor=0, post_translate=True
+        self,
+        R,
+        t=None,
+        pose_rotvecs=None,
+        shape_betas=None,
+        trans=None,
+        kid_factor=0,
+        post_translate=True,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         Rotates and translates the body in parametric form.
@@ -239,7 +246,7 @@ def rototranslate(
 
         Parameters:
             R: Rotation matrix, shaped as (3, 3).
-            t: Translation vector, shaped as (3,).
+            t: Translation vector, shaped as (3,). Defaults to zero (pure rotation).
             pose_rotvecs: Initial rotation vectors per joint, shaped as (num_joints * 3,).
             shape_betas: Shape coefficients (betas) for body shape, shaped as (num_betas,).
             trans: Initial translation vector, shaped as (3,).
@@ -261,6 +268,10 @@ def rototranslate(
             account the offset between the pelvis joint in the shaped T-pose and the origin of
             the canonical coordinate system.
         """
+        if t is None:
+            t = np.zeros(3, dtype=R.dtype)
+        if pose_rotvecs is None or shape_betas is None or trans is None:
+            raise ValueError('pose_rotvecs, shape_betas, and trans are required.')
         current_rotmat = rotvec2mat(pose_rotvecs[:3])
         new_rotmat = R @ current_rotmat
         new_pose_rotvec = np.concatenate([mat2rotvec(new_rotmat), pose_rotvecs[3:]], axis=0)

src/smplfitter/np/__init__.py

@@ -8,7 +8,6 @@
 from smplfitter.common import _set_module_for_docs
 
 import functools
-import os
 
 __all__ = ['BodyModel', 'BodyFitter', 'BodyConverter', 'get_cached_body_model']
 _set_module_for_docs(__name__, globals(), __all__)
@@ -20,7 +19,4 @@ def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
 
 
 def get_body_model(model_name, gender, model_root=None):
-    if model_root is None:
-        DATA_ROOT = os.getenv('DATA_ROOT', default='.')
-        model_root = f'{DATA_ROOT}/body_models/{model_name}'
     return BodyModel(model_root=model_root, gender=gender, model_name=model_name)

src/smplfitter/np/bodymodel.py

@@ -232,7 +232,14 @@ def single(self, *args, return_vertices=True, **kwargs):
         return {k: np.squeeze(v, axis=0) for k, v in result.items()}
 
     def rototranslate(
-        self, R, t, pose_rotvecs, shape_betas, trans, kid_factor=0, post_translate=True
+        self,
+        R,
+        t=None,
+        pose_rotvecs=None,
+        shape_betas=None,
+        trans=None,
+        kid_factor=0,
+        post_translate=True,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         Rotates and translates the body in parametric form.
@@ -248,7 +255,7 @@ def rototranslate(
 
         Parameters:
             R: Rotation matrix, shaped as (3, 3).
-            t: Translation vector, shaped as (3,).
+            t: Translation vector, shaped as (3,). Defaults to zero (pure rotation).
             pose_rotvecs: Initial rotation vectors per joint, shaped as (num_joints * 3,).
             shape_betas: Shape coefficients (betas) for body shape, shaped as (num_betas,).
             trans: Initial translation vector, shaped as (3,).
@@ -270,6 +277,10 @@ def rototranslate(
             account the offset between the pelvis joint in the shaped T-pose and the origin of
             the canonical coordinate system.
         """
+        if t is None:
+            t = np.zeros(3, dtype=R.dtype)
+        if pose_rotvecs is None or shape_betas is None or trans is None:
+            raise ValueError('pose_rotvecs, shape_betas, and trans are required.')
         current_rotmat = rotvec2mat(pose_rotvecs[:3])
         new_rotmat = R @ current_rotmat
         new_pose_rotvec = np.concatenate([mat2rotvec(new_rotmat), pose_rotvecs[3:]], axis=0)

src/smplfitter/pt/__init__.py

@@ -5,7 +5,6 @@
 from typing import Optional
 
 import functools
-import os
 import torch
 import warnings
 from .bodymodel import BodyModel
@@ -33,9 +32,6 @@ def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
 
 
 def get_body_model(model_name, gender, model_root=None):
-    if model_root is None:
-        DATA_ROOT = os.getenv('DATA_ROOT', default='.')
-        model_root = f'{DATA_ROOT}/body_models/{model_name}'
     return BodyModel(model_root=model_root, gender=gender, model_name=model_name)
 
 

src/smplfitter/pt/bodymodel.py

@@ -338,10 +338,10 @@ def single(
     def rototranslate(
         self,
         R: torch.Tensor,
-        t: torch.Tensor,
-        pose_rotvecs: torch.Tensor,
-        shape_betas: torch.Tensor,
-        trans: torch.Tensor,
+        t: Optional[torch.Tensor] = None,
+        pose_rotvecs: Optional[torch.Tensor] = None,
+        shape_betas: Optional[torch.Tensor] = None,
+        trans: Optional[torch.Tensor] = None,
         kid_factor: Optional[torch.Tensor] = None,
         post_translate: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -387,6 +387,10 @@ def rototranslate(
 
         """
 
+        if t is None:
+            t = torch.zeros(3, device=R.device, dtype=R.dtype)
+        if pose_rotvecs is None or shape_betas is None or trans is None:
+            raise ValueError('pose_rotvecs, shape_betas, and trans are required.')
         current_rotmat = rotvec2mat(pose_rotvecs[:3])
         new_rotmat = R @ current_rotmat
         new_pose_rotvec = torch.cat([mat2rotvec(new_rotmat), pose_rotvecs[3:]], dim=0)

src/smplfitter/tf/__init__.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import functools
-import os
 
 import tensorflow as tf
 from .bodymodel import BodyModel
@@ -27,9 +26,6 @@ def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
 
 
 def get_body_model(model_name, gender, model_root=None, num_betas=None, vertex_subset=None):
-    if model_root is None:
-        DATA_ROOT = os.getenv('DATA_ROOT', '.')
-        model_root = f'{DATA_ROOT}/body_models/{model_name}'
     return BodyModel(
         model_root=model_root,
         gender=gender,

src/smplfitter/tf/bodymodel.py

@@ -249,10 +249,10 @@ def single(self, *args, return_vertices=True, **kwargs):
     def rototranslate(
         self,
         R: tf.Tensor,
-        t: tf.Tensor,
-        pose_rotvecs: tf.Tensor,
-        shape_betas: tf.Tensor,
-        trans: tf.Tensor,
+        t=None,
+        pose_rotvecs=None,
+        shape_betas=None,
+        trans=None,
         kid_factor=0,
         post_translate: bool = True,
     ):
@@ -291,6 +291,10 @@ def rototranslate(
             account the offset between the pelvis joint in the shaped T-pose and the origin of
             the canonical coordinate system.
         """
+        if t is None:
+            t = tf.zeros(3, dtype=R.dtype)
+        if pose_rotvecs is None or shape_betas is None or trans is None:
+            raise ValueError('pose_rotvecs, shape_betas, and trans are required.')
         current_rotmat = rotvec2mat(pose_rotvecs[..., :3])
         new_rotmat = R @ current_rotmat
         new_pose_rotvec = tf.concat([mat2rotvec(new_rotmat), pose_rotvecs[3:]], axis=0)