Add a method to support direct SAM3D-Body output to SMPL(X) (#43)

Co-authored-by: Jinlong Yang <jinlong.yang@oculus.com>

Author: JinlongYANG

tools/mhr_smpl_conversion/README.md

@@ -8,6 +8,7 @@ The SMPL-MHR Conversion Tool enables seamless conversion between different 3D hu
 
 - **SMPL/SMPLX → MHR**: Convert SMPL/SMPLX model parameters/meshes to MHR format
 - **MHR → SMPL/SMPLX**: Convert MHR parameters back to SMPL/SMPLX format
+- **SAM3D Outputs (MHR) → SMPL/SMPLX**: Convert SAM3D output to SMPL/SMPLX model parameters
 
 The tool uses barycentric interpolation for topology mapping and offers multiple optimization backends for parameter fitting.
 
@@ -64,6 +65,7 @@ mhr_smpl_conversion/
 │   ├── smpl_para2mhr_pymomentum
 │   ├── mhr_para2smplx_pytorch_single_identity
 │   ├── smplx_mesh2mhr_pytorch_single_identity
+│   ├── sam3d_output_to_smplx
 └──...
 ```
 
@@ -88,6 +90,14 @@ converter = Conversion(
     method="pytorch"  # or "pymomentum"
 )
 
+# Convert SAM3D output to SMPL(X)
+results = converter.convert_sam3d_output_to_smplx(
+    sam3d_outputs=sam3d_outputs,
+    return_smpl_meshes=True,
+    return_smpl_parameters=True,
+    return_fitting_errors=True
+)
+
 # Convert SMPLX to MHR
 results = converter.convert_smpl2mhr(
     smpl_parameters=smplx_params,

tools/mhr_smpl_conversion/conversion.py

@@ -82,7 +82,7 @@
 
 import logging
 from functools import cached_property, lru_cache
-from typing import Optional
+from typing import Optional, Any
 
 import numpy as np
 
@@ -322,9 +322,9 @@ def convert_smpl2mhr(
         return ConversionResult(
             result_meshes=result_meshes if return_mhr_meshes else None,
             result_vertices=mhr_vertices if return_mhr_vertices else None,
-            result_parameters=None
-            if not return_mhr_parameters
-            else fitting_parameter_results,
+            result_parameters=(
+                None if not return_mhr_parameters else fitting_parameter_results
+            ),
             result_errors=errors if return_fitting_errors else None,
         )
 
@@ -401,12 +401,132 @@ def convert_mhr2smpl(
         return ConversionResult(
             result_meshes=result_meshes if return_smpl_meshes else None,
             result_vertices=smpl_vertices if return_smpl_vertices else None,
-            result_parameters=fitting_parameter_results
-            if return_smpl_parameters
-            else None,
+            result_parameters=(
+                fitting_parameter_results if return_smpl_parameters else None
+            ),
             result_errors=errors if return_fitting_errors else None,
         )
 
+    def convert_sam3d_output_to_smpl(
+        self,
+        sam3d_outputs: list[dict[str, Any]],
+        return_smpl_meshes: bool = False,
+        return_smpl_parameters: bool = True,
+        return_smpl_vertices: bool = False,
+        return_fitting_errors: bool = True,
+        batch_size: int = 256,
+    ) -> ConversionResult:
+        """
+        Convert SAM3D output to SMPL model parameters.
+
+        Usage:
+            After initializing a coverter object from Conversion class,
+            sam3d_outputs = SAM3DBodyEstimator.process_one_image(...)
+            conversion_results = coverter.convert_sam3d_output_to_smpl(sam3d_outputs)
+
+        Args:
+            sam3d_outputs: List of dictionaries containing the output of the SAM3D model.
+            return_smpl_meshes: Whether to return the SMPL meshes. If True, the function will return a list
+                                of SMPL meshes.
+            return_smpl_parameters: Whether to return the SMPL parameters. If True, the function will return a
+                                    dictionary of SMPL parameters.
+            return_smpl_vertices: Whether to return the SMPL vertices. If True, the function will return a numpy
+                                  array of SMPL vertices.
+            return_fitting_errors: Whether to return fitting errors. If True, the function will return errors for
+                                   each frame.
+            batch_size: Number of frames to process in each batch.
+
+        Returns:
+            ConversionResult containing:
+                - result_meshes: List of SMPL meshes
+                - result_vertices: Numpy array of SMPL vertices
+                - result_parameters: Dictionary of SMPL parameters
+                - result_errors: Numpy array of fitting errors for each frame
+        """
+
+        # These are the keys that are expected in the sam3d_outputs, from the first three we can extract the MHR vertices. The last one are the MHR vertices.
+        SAM3D_OUTPUT_VARIABLES = (
+            "mhr_model_params",  # Translation is not included in the model params, but in "pred_cam_t".
+            "shape_params",
+            "expr_params",
+            "pred_vertices",  # Translation is not included in the vertices, but in "pred_cam_t".
+            "pred_cam_t",  # This is the camera translation w.r.t. each person.
+        )
+
+        # The result of SAM3DBodyEstimator.process_one_image() is a list of dictionaries, each dictionary contains the outputs of the SAM3D model for one person.
+        # We concatenate all person's outputs into one dictionary, so that we can convert them to SMPL in batch.
+        concatenated_sam3d_outputs = {}
+        for k in SAM3D_OUTPUT_VARIABLES:
+            concatenated_sam3d_outputs[k] = np.stack(
+                [
+                    sam3d_output[k]
+                    for sam3d_output in sam3d_outputs
+                    if k in sam3d_output
+                ],
+                axis=0,
+            )
+            concatenated_sam3d_outputs[k] = self._to_tensor(
+                concatenated_sam3d_outputs[k]
+            )
+
+        # If MHR vertices are available, use them to compute the SMPL parameters. If not, use the MHR model, shape, and expressions to compute the SMPL parameters.
+        num_people = len(sam3d_outputs)
+        mhr_vertices = None
+        if (
+            "pred_vertices" in concatenated_sam3d_outputs
+            and concatenated_sam3d_outputs["pred_vertices"].shape[0] == num_people
+        ):
+            # If pred_vertices is available, use it to compute the target vertices
+            # Note: SAM3D outputs the vertices in meters, so we need to convert them to centimeters before passing them to the conversion function
+            mhr_vertices = (
+                100.0 * concatenated_sam3d_outputs["pred_vertices"]
+                + 100.0 * concatenated_sam3d_outputs["pred_cam_t"][:, None, :]
+            )
+        elif (
+            "mhr_model_params" in concatenated_sam3d_outputs
+            and concatenated_sam3d_outputs["mhr_model_params"].shape[0] == num_people
+        ):
+            # If pred_vertices is not available, use the mhr_model_params to compute the target vertices
+            mhr_parameters = {}
+            mhr_parameters["lbs_model_params"] = concatenated_sam3d_outputs[
+                "mhr_model_params"
+            ]
+
+            assert (
+                "shape_params" in concatenated_sam3d_outputs
+                and concatenated_sam3d_outputs["shape_params"].shape[0] == num_people
+            )
+            mhr_parameters["identity_coeffs"] = concatenated_sam3d_outputs[
+                "shape_params"
+            ]
+            assert (
+                "expr_params" in concatenated_sam3d_outputs
+                and concatenated_sam3d_outputs["expr_params"].shape[0] == num_people
+            )
+            mhr_parameters["face_expr_coeffs"] = concatenated_sam3d_outputs[
+                "expr_params"
+            ]
+            _, mhr_vertices = self._mhr_para2mesh(mhr_parameters, return_mesh=False)
+            mhr_vertices = self._to_tensor(mhr_vertices)
+            mhr_vertices[..., [1, 2]] *= -1  # Camera system difference in SAM3D-Body
+            mhr_vertices += 100.0 * concatenated_sam3d_outputs["pred_cam_t"][:, None, :]
+        else:
+            raise ValueError(
+                "Either pred_vertices or mhr_model_params must be available in the SAM3D output."
+            )
+
+        conversion_results = self.convert_mhr2smpl(
+            mhr_vertices=mhr_vertices,
+            single_identity=False,
+            is_tracking=False,
+            return_smpl_meshes=return_smpl_meshes,
+            return_smpl_parameters=return_smpl_parameters,
+            return_smpl_vertices=return_smpl_vertices,
+            return_fitting_errors=return_fitting_errors,
+            batch_size=batch_size,
+        )
+        return conversion_results
+
     def _to_tensor(self, data: torch.Tensor | np.ndarray) -> torch.Tensor:
         """Convert input data to tensor on the appropriate device."""
         if isinstance(data, torch.Tensor):

tools/mhr_smpl_conversion/data/sam3d_body_outputs/img_subj00.npz

@@ -29,6 +29,9 @@
 from conversion import Conversion
 
 _INPUT_FILE = "./data/example_smplx_poses.npy"  # Directory to store input data
+_SAM3D_BODY_OUTPUT_DIR = (
+    "./data/sam3d_body_outputs"  # Directory to store SAM3D body output
+)
 _OUTPUT_DIR = "./tmp_results"  # Directory to store conversion results
 
 
@@ -67,6 +70,7 @@ def run_examples(
                     if not os.path.exists(converted_smpl_model_file):
                         smpl_model_data = dict(np.load(smpl_model_file))
                         import pickle
+
                         with open(converted_smpl_model_file, "wb") as f:
                             pickle.dump(smpl_model_data, f)
                     self.smpl_model = smplx.SMPL(
@@ -97,12 +101,8 @@ def run_examples(
 
         # Example 2: Convert the SMPLX meshes to MHR with PyTorch and single identity
         if self.smplx_model is not None:
-            print(
-                "\nConverting SMPLX meshes to MHR with PyTorch and single identity."
-            )
-            example_output_dir = (
-                output_dir + "/smplx_mesh2mhr_pytorch_single_identity"
-            )
+            print("\nConverting SMPLX meshes to MHR with PyTorch and single identity.")
+            example_output_dir = output_dir + "/smplx_mesh2mhr_pytorch_single_identity"
             os.makedirs(example_output_dir, exist_ok=True)
             mhr_parameters = (
                 self.example_smplx_meshes_to_mhr_with_pytorch_single_identity(
@@ -119,6 +119,54 @@ def run_examples(
                 mhr_parameters, example_output_dir
             )
 
+        # Example 4: Convert the SAM3D output to SMPLX
+        if self.smplx_model is not None:
+            print("\nConverting SAM3D output to SMPLX.")
+            example_output_dir = output_dir + "/sam3d_output_to_smplx"
+            os.makedirs(example_output_dir, exist_ok=True)
+            self.example_sam3d_output_to_smplx(example_output_dir)
+
+    def example_sam3d_output_to_smplx(self, output_dir: str):
+        """Convert SAM3D output to SMPLX."""
+        mhr_model = MHR.from_files(lod=1, device=self._device)
+
+        # Load the SAM3D outputs
+        subject_files = [
+            f for f in os.listdir(_SAM3D_BODY_OUTPUT_DIR) if f.endswith(".npz")
+        ]
+        sam3d_outputs = []
+
+        for subject_file in subject_files:
+            subject_data = np.load(os.path.join(_SAM3D_BODY_OUTPUT_DIR, subject_file))
+            sam3d_outputs.append(subject_data)
+            mesh = trimesh.Trimesh(
+                subject_data["pred_vertices"] + subject_data["pred_cam_t"][None, ...],
+                mhr_model.character.mesh.faces,
+                process=False,
+            )
+            mesh.export(f"{output_dir}/{subject_file[:-4]}_sam3d.ply")
+
+        # ***** Core conversion code *****
+        converter = Conversion(
+            mhr_model=mhr_model, smpl_model=self.smplx_model, method="pytorch"
+        )
+
+        conversion_results = converter.convert_sam3d_output_to_smpl(
+            sam3d_outputs=sam3d_outputs,
+            return_smpl_meshes=True,
+            return_smpl_parameters=False,
+            return_smpl_vertices=False,
+            return_fitting_errors=True,
+        )
+
+        print("Conversion errors:")
+        print(conversion_results.result_errors)
+
+        # Save the results
+        print(f"Exporting the results to {output_dir}...")
+        for i, mesh in enumerate(conversion_results.result_meshes):
+            mesh.export(f"{output_dir}/{subject_files[i][:-4]}_result_smplx.ply")
+
     def example_smpl_parameters_to_mhr_with_pymomentum_multiple_identity(
         self, smpl_paras: dict[str, torch.Tensor], output_dir: str
     ):
@@ -301,7 +349,7 @@ def _parse_arguments():
         "--smpl",
         type=str,
         required=False,
-        default="./data/SMPL_NEUTRAL.npz",
+        default="./data/SMPL_NEUTRAL.pkl",
         help="Path to the SMPL model file. If not specified, SMPL related conversion will be skipped.",
     )