From c787361f2583fb710475e48d9387357d1d32846e Mon Sep 17 00:00:00 2001
From: mkocabas <muhammed.kocabas@metu.edu.tr>
Date: Tue, 17 Dec 2019 16:46:12 +0100
Subject: [PATCH] option to save output meshes as obj files

---
 demo.py               | 11 +++++++++++
 doc/demo.md           |  2 ++
 lib/utils/renderer.py |  5 ++++-
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/demo.py b/demo.py
index b7ef073..e948cba 100644
--- a/demo.py
+++ b/demo.py
@@ -277,11 +277,19 @@ def main(args):
 
                 mc = mesh_color[person_id]
 
+                mesh_filename = None
+
+                if args.save_obj:
+                    mesh_folder = os.path.join(output_path, 'meshes', f'{person_id:04d}')
+                    os.makedirs(mesh_folder, exist_ok=True)
+                    mesh_filename = os.path.join(mesh_folder, f'{frame_idx:06d}.obj')
+
                 img = renderer.render(
                     img,
                     frame_verts,
                     cam=frame_cam,
                     color=mc,
+                    mesh_filename=mesh_filename,
                 )
 
                 if args.sideview:
@@ -361,6 +369,9 @@ def main(args):
     parser.add_argument('--sideview', action='store_true',
                         help='render meshes from alternate viewpoint.')
 
+    parser.add_argument('--save_obj', action='store_true',
+                        help='save results as .obj files.')
+
     args = parser.parse_args()
 
     main(args)
diff --git a/doc/demo.md b/doc/demo.md
index 0fd4bef..6245931 100644
--- a/doc/demo.md
+++ b/doc/demo.md
@@ -36,6 +36,8 @@ For this option, you have to set `--tracking_method` option to `pose`.
 - `--sideview`: Render the output meshes from an alternate viewpoint. Default alternate viewpoint is -90 degrees in y axis.
 Note that this option doubles the rendering time.
 
+- `--save_obj`: Save output meshes as .obj files.
+
 ## Examples
 - Run VIBE on a video file using bbox tracker and visualize the results with wireframe meshes:
 ```bash
diff --git a/lib/utils/renderer.py b/lib/utils/renderer.py
index 549d5ea..63cc228 100644
--- a/lib/utils/renderer.py
+++ b/lib/utils/renderer.py
@@ -75,13 +75,16 @@ def __init__(self, resolution=(224,224), orig_img=False, wireframe=False):
         light_pose[:3, 3] = [1, 1, 2]
         self.scene.add(light, pose=light_pose)
 
-    def render(self, img, verts, cam, angle=None, axis=None, color=[1.0, 1.0, 0.9]):
+    def render(self, img, verts, cam, angle=None, axis=None, mesh_filename=None, color=[1.0, 1.0, 0.9]):
 
         mesh = trimesh.Trimesh(vertices=verts, faces=self.faces)
 
         Rx = trimesh.transformations.rotation_matrix(math.radians(180), [1, 0, 0])
         mesh.apply_transform(Rx)
 
+        if mesh_filename is not None:
+            mesh.export(mesh_filename)
+
         if angle and axis:
             R = trimesh.transformations.rotation_matrix(math.radians(angle), axis)
             mesh.apply_transform(R)