Add option to save output meshes.

README.md

@@ -4,7 +4,7 @@ Code repository for the paper:
 **Score-Guided Diffusion for 3D Human Recovery**\
 [Anastasis Stathopoulos](https://statho.github.io/), [Ligong Han](https://phymhan.github.io/), [Dimitris Metaxas](https://people.cs.rutgers.edu/~dnm/)
 
-[![arXiv](https://img.shields.io/badge/arXiv-2305.20091-00ff00.svg)](https://arxiv.org/abs/2403.09623)  [![Website shields.io](https://img.shields.io/website-up-down-green-red/http/shields.io.svg)](https://statho.github.io/ScoreHMR/)  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1WdEWaOyG0XHHZrlg2uX9XwECDKBX5dCn?usp=sharing)
+[![arXiv](https://img.shields.io/badge/arXiv-2305.20091-00ff00.svg)](https://arxiv.org/abs/2403.09623)  [![Website shields.io](https://img.shields.io/website-up-down-green-red/http/shields.io.svg)](https://statho.github.io/ScoreHMR/)  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1WdEWaOyG0XHHZrlg2uX9XwECDKBX5dCn)
 
 ![teaser](assets/teaser.jpg)
 
@@ -36,7 +36,7 @@ python demo_image.py \
 
 
 ## Run demo on videos
-The following command will first run tracking with 4D-Humans and 2D keypoint detection with ViTPose, and then run temporal model fitting with ScoreHMR on the video specified with `--input_video`. It will create a video rendering of the reconstructed people in the folder specified by `--output_folder`. It will also save intermediate results from 4D-Humans and ViTPose.
+The following command will first run tracking with 4D-Humans and 2D keypoint detection with ViTPose, and then run temporal model fitting with ScoreHMR on the video specified with `--input_video`. It will create a video rendering of the reconstructed people in the folder specified by `--out_folder`. It will also save intermediate results from 4D-Humans and ViTPose.
 ```bash
 python demo_video.py \
     --input_video example_data/videos/breakdancing.mp4 \

demo_video.py

@@ -16,14 +16,16 @@
 warnings.filterwarnings('ignore')
 
 
+LIGHT_BLUE=(0.65098039,  0.74117647,  0.85882353)
 NUM_SAMPLES = 1
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input_video", type=str, default="example_data/breakdancing.mp4", help="Path of the input video.")
     parser.add_argument("--out_folder", type=str, default="demo_out/videos", help="Path to save the output video.")
-    parser.add_argument("--overwrite", action="store_true", default=False, help="Indicate whether or not to overwrite the 4D-Human tracklets.")
+    parser.add_argument("--overwrite", action="store_true", default=False, help="If set, overwrite the 4D-Human tracklets.")
+    parser.add_argument('--save_mesh', action='store_true', default=False, help='If set, save meshes to disk.')
     parser.add_argument("--fps", type=int, default=30, help="Frame rate to save the output video.")
     args = parser.parse_args()
 
@@ -165,6 +167,19 @@ def main():
             pred_cam_t_all[track_idx, start_idx:end_idx] = dm_out['camera_translation'].cpu()
             pred_vertices_all[track_idx, start_idx:end_idx] = smpl_out.vertices.cpu()
 
+            # Save meshes as OBJ files.
+            if args.save_mesh:
+                verts = smpl_out.vertices.cpu().numpy()
+                cam_t = dm_out['camera_translation'].cpu().numpy()
+                person_id = str(batch['track_id'].item()).zfill(3)
+                tmesh_path = f"{OUT_DIR}/mesh_output/{filename}/{person_id}"
+                print(f'=> Saving mesh files for {person_id} in {tmesh_path}')
+                os.makedirs(tmesh_path, exist_ok=True)
+                for ii, (vvv, ttt) in enumerate(zip(verts, cam_t)):
+                    tmesh = renderer.vertices_to_trimesh(vvv, ttt, LIGHT_BLUE)
+                    frame_id = str(ii + start_idx + 1).zfill(6)
+                    tmesh.export(f'{tmesh_path}/{frame_id}.obj')
+
 
         # Save output video.
         frame_list = create_visuals(