Update render

app.py

@@ -6,6 +6,7 @@ import cv2
 import os
 import numpy as np
 import OpenGL.GL as gl
+import imageio
 import pytorch_lightning as pl
 import moviepy.editor as mp
 from pathlib import Path
@@ -118,33 +119,26 @@ def render_motion(data, feats, method='fast'):
         shape = [768, 768]
         render = SMPLRender(cfg.RENDER.SMPL_MODEL_PATH)
 
-        if not os.environ.get("PYOPENGL_PLATFORM"):
-            os.environ["DISPLAY"] = ":0.0"
-            os.environ["PYOPENGL_PLATFORM"] = "egl"
-
-        size = (shape[1], shape[0])
-        fps = 20.0
-        fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')
-        videoWriter = cv2.VideoWriter(output_mp4_path, fourcc, fps, size)
         r = RRR.from_rotvec(np.array([np.pi, 0.0, 0.0]))
         pose[:, 0] = np.matmul(r.as_matrix().reshape(1, 3, 3), pose[:, 0])
+        vid = []
+        aroot = data[[0], 0]
+        aroot[:, 1] = -aroot[:, 1]
+        params = dict(pred_shape=np.zeros([1, 10]),
+                        pred_root=aroot,
+                        pred_pose=pose)
+        render.init_renderer([shape[0], shape[1], 3], params)
         for i in range(data.shape[0]):
-            img = np.zeros([shape[0], shape[1], 3])
-            aroot = data[[i], 0] + np.array([[0.0, 0.0, 30.0]])
-            aroot[:, 1] = -aroot[:, 1]
-            params = dict(pred_shape=np.zeros([1, 10]),
-                          pred_root=aroot,
-                          pred_pose=pose[[i]])
-            renderImg = render.render(img.copy(), params)
-            renderImg = (renderImg * 255).astype(np.uint8)
-            videoWriter.write(renderImg)
-        videoWriter.release()
-        output_video_h264_name = output_mp4_path[:-4] + '_h264.mp4'
-        command = 'ffmpeg -y -i {} -vcodec h264 {}'.format(
-            output_mp4_path, output_video_h264_name)
-        os.system(command)
-        output_mp4_path = output_video_h264_name
-        video_fname = video_fname[:-4] + '_h264.mp4'
+            renderImg = render.render(i)
+            vid.append(renderImg)
+            
+        out = np.stack(vid, axis=0)
+        output_gif_path = output_mp4_path[:-4] + '.gif'
+        imageio.mimwrite(output_gif_path, out, duration=50)
+        out_video = mp.VideoFileClip(output_gif_path)
+        out_video.write_videofile(output_mp4_path)
+        del out, render
+        
     elif method == 'fast':
         output_gif_path = output_mp4_path[:-4] + '.gif'
         if len(data.shape) == 3:
@@ -154,6 +148,7 @@ def render_motion(data, feats, method='fast'):
         pose_vis = plot_3d.draw_to_batch(data, [''], [output_gif_path])
         out_video = mp.VideoFileClip(output_gif_path)
         out_video.write_videofile(output_mp4_path)
+        del pose_vis
 
     return output_mp4_path, video_fname, output_npy_path, feats_fname
 
@@ -544,7 +539,7 @@ with gr.Blocks(css=customCSS) as demo:
                                  label="Visulization method",
                                  interactive=True,
                                  elem_id="method",
-                                 value="fast")
+                                 value="slow")
 
             language = gr.Dropdown(["English", "中文"],
                                    label="Speech language",

mGPT/render/pyrender/smpl_render.py

@@ -55,7 +55,7 @@ class Renderer:
                  faces=None):
         self.renderer = pyrender.OffscreenRenderer(viewport_width=img_res[0],
                                                    viewport_height=img_res[1],
-                                                   point_size=1.0)
+                                                   point_size=2.0)
         self.focal_length = focal_length
         self.camera_center = [img_res[0] // 2, img_res[1] // 2]
         self.faces = faces
@@ -65,58 +65,62 @@ class Renderer:
         else:
             self.device = torch.device("cpu")
 
-        vertices = np.concatenate(vertices)
-        # Center the first root to the first frame
-        vertices -= vertices[[0], [0], :]
-        # Remove the floor
-        vertices[..., 2] -= vertices[..., 2].min()
-        data = vertices[..., [2, 0, 1]]
-        minx, miny, _ = data.min(axis=(0, 1))
-        maxx, maxy, _ = data.max(axis=(0, 1))
-        minz, maxz = -0.5, 0.5
-        minx = minx - 0.5
-        maxx = maxx + 0.5
-        miny = miny - 0.5
-        maxy = maxy + 0.5
-
-        polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz],
-                                    [maxx, minz]])
-        self.polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)
-        self.polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]
         self.rot = trimesh.transformations.rotation_matrix(
             np.radians(180), [1, 0, 0])
-        # self.polygon_mesh.apply_transform(self.rot)
+
+        minx, miny, minz = vertices.min(axis=(0, 1))
+        maxx, maxy, maxz = vertices.max(axis=(0, 1))
+        minx = minx - 0.5
+        maxx = maxx + 0.5
+        minz = minz - 0.5
+        maxz = maxz + 0.5
+
+        floor = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz],
+                                  [maxx, minz]])
+        self.floor = trimesh.creation.extrude_polygon(floor, 1e-5)
+        self.floor.visual.face_colors = [0, 0, 0, 0.2]
+        self.floor.apply_transform(self.rot)
+        self.floor_pose = np.array(
+            [[1, 0, 0, 0], [0, np.cos(np.pi / 2), -np.sin(np.pi / 2), miny],
+             [0, np.sin(np.pi / 2), np.cos(np.pi / 2), 0], [0, 0, 0, 1]])
+
+        c = -np.pi / 6
+        self.camera_pose = [[1, 0, 0, (minx + maxx) / 2],
+                            [0, np.cos(c), -np.sin(c), 1.5],
+                            [
+                                0,
+                                np.sin(c),
+                                np.cos(c),
+                                max(4, minz + (1.5 - miny) * 2, (maxx - minx))
+                            ], [0, 0, 0, 1]]
 
     def __call__(self, vertices, camera_translation):
-        scene = pyrender.Scene(bg_color=(1., 1., 1., 0.8),
-                               ambient_light=(0.4, 0.4, 0.4))
+
+        floor_render = pyrender.Mesh.from_trimesh(self.floor, smooth=False)
 
         material = pyrender.MetallicRoughnessMaterial(
-            metallicFactor=0.4,
+            metallicFactor=0.1,
             alphaMode='OPAQUE',
             baseColorFactor=(0.658, 0.214, 0.0114, 0.2))
         mesh = trimesh.Trimesh(vertices, self.faces)
         mesh.apply_transform(self.rot)
         mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
+
+        camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0), znear=0.5)
+
+        light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=350)
+        spot_l = pyrender.SpotLight(color=np.ones(3),
+                                    intensity=300.0,
+                                    innerConeAngle=np.pi / 16,
+                                    outerConeAngle=np.pi / 6)
+        point_l = pyrender.PointLight(color=np.ones(3), intensity=300.0)
+
+        scene = pyrender.Scene(bg_color=(1., 1., 1., 0.8),
+                               ambient_light=(0.4, 0.4, 0.4))
+        scene.add(floor_render, pose=self.floor_pose)
         scene.add(mesh, 'mesh')
 
-        polygon_render = pyrender.Mesh.from_trimesh(self.polygon_mesh,
-                                                    smooth=False)
-        c = np.pi / 2
-        scene.add(polygon_render)
-
-        camera_pose = np.eye(4)
-        camera_translation[0] *= -1.
-        camera_pose[:3, 3] = camera_translation
-        camera = pyrender.IntrinsicsCamera(fx=self.focal_length,
-                                           fy=self.focal_length,
-                                           cx=self.camera_center[0],
-                                           cy=self.camera_center[1])
-        scene.add(camera, pose=camera_pose)
-
-        light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=300)
         light_pose = np.eye(4)
-
         light_pose[:3, 3] = np.array([0, -1, 1])
         scene.add(light, pose=light_pose)
 
@@ -126,8 +130,10 @@ class Renderer:
         light_pose[:3, 3] = np.array([1, 1, 2])
         scene.add(light, pose=light_pose)
 
-        color, rend_depth = self.renderer.render(
-            scene, flags=pyrender.RenderFlags.RGBA)
+        scene.add(camera, pose=self.camera_pose)
+
+        flags = pyrender.RenderFlags.RGBA | pyrender.RenderFlags.SHADOWS_DIRECTIONAL
+        color, rend_depth = self.renderer.render(scene, flags=flags)
 
         return color
 
@@ -141,27 +147,35 @@ class SMPLRender():
             self.device = torch.device("cpu")
         self.smpl = SMPL(SMPL_MODEL_DIR, batch_size=1,
                          create_transl=False).to(self.device)
-        self.vertices = []
+
         self.pred_camera_t = []
-        self.focal_length = 5000
-
-    def fit(self, smpl_param, is_headroot=False):
-        pose = smpl_param['pred_pose']
-        if pose.size == 72:
-            pose = pose.reshape(-1, 3)
-            pose = RRR.from_rotvec(pose).as_matrix()
-            pose = pose.reshape(1, 24, 3, 3)
+        self.focal_length = 110
+
+    def init_renderer(self, res, smpl_param, is_headroot=False):
+        poses = smpl_param['pred_pose']
+        pred_rotmats = []
+        for pose in poses:
+            if pose.size == 72:
+                pose = pose.reshape(-1, 3)
+                pose = RRR.from_rotvec(pose).as_matrix()
+                pose = pose.reshape(1, 24, 3, 3)
+            pred_rotmats.append(
+                torch.from_numpy(pose.astype(np.float32)[None]).to(
+                    self.device))
+
+        pred_rotmat = torch.cat(pred_rotmats, dim=0)
+
         pred_betas = torch.from_numpy(smpl_param['pred_shape'].reshape(
             1, 10).astype(np.float32)).to(self.device)
-        pred_rotmat = torch.from_numpy(pose.astype(np.float32)).to(self.device)
         pred_camera_t = smpl_param['pred_root'].reshape(1,
                                                         3).astype(np.float32)
+
         smpl_output = self.smpl(betas=pred_betas,
                                 body_pose=pred_rotmat[:, 1:],
                                 global_orient=pred_rotmat[:, 0].unsqueeze(1),
                                 pose2rot=False)
-        vertices = smpl_output.vertices[0].detach().cpu().numpy()
-        self.vertices.append(vertices[None])
+
+        self.vertices = smpl_output.vertices.detach().cpu().numpy()
 
         pred_camera_t = pred_camera_t[0]
 
@@ -171,13 +185,12 @@ class SMPLRender():
 
         self.pred_camera_t.append(pred_camera_t)
 
-    def init_renderer(self, res):
         self.renderer = Renderer(vertices=self.vertices,
                                  focal_length=self.focal_length,
                                  img_res=(res[1], res[0]),
                                  faces=self.smpl.faces)
 
     def render(self, index):
-        renderImg = self.renderer(self.vertices[index][0],
-                                  self.pred_camera_t[index].copy())
+        renderImg = self.renderer(self.vertices[index, ...],
+                                  self.pred_camera_t)
         return renderImg