Update render method

.gitignore

@@ -8,7 +8,6 @@ __pycache__/
 .DS_Store
 pyglet
 app2.py
-render.py
 cache
 
 # Distribution / packaging

app.py

@@ -7,27 +7,18 @@ os.system('pip install pyglet==1.4.0a1')
 os.system('pip install triangle==20220202')
 
 import gradio as gr
-import random
 import torch
 import time
-import cv2
 import numpy as np
-import OpenGL.GL as gl
-import imageio
 import pytorch_lightning as pl
-import moviepy.editor as mp
+import subprocess
 from pathlib import Path
 from mGPT.data.build_data import build_data
 from mGPT.models.build_model import build_model
 from mGPT.config import parse_args
-from scipy.spatial.transform import Rotation as RRR
-import mGPT.render.matplot.plot_3d_global as plot_3d
-from mGPT.render.pyrender.hybrik_loc2rot import HybrIKJointsToRotmat
-from mGPT.render.pyrender.smpl_render import SMPLRender
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import librosa
 from huggingface_hub import snapshot_download
-import eventlet
 
 # Load model
 cfg = parse_args(phase="webui")  # parse config file
@@ -105,55 +96,17 @@ def render_motion(data, feats, method='fast'):
     fname = time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime(
         time.time())) + str(np.random.randint(10000, 99999))
     video_fname = fname + '.mp4'
-    feats_fname = fname + '.npy'
+    feats_fname = f"{fname}_feats" + '.npy'
+    data_fname = f"{fname}_joints" + '.npy'
     output_npy_path = os.path.join(output_dir, feats_fname)
+    output_joints_path = os.path.join(output_dir, data_fname)
     output_mp4_path = os.path.join(output_dir, video_fname)
     np.save(output_npy_path, feats)
-
-    if method == 'slow':
-        if len(data.shape) == 4:
-            data = data[0]
-        data = data - data[0, 0]
-        pose_generator = HybrIKJointsToRotmat()
-        pose = pose_generator(data)
-        pose = np.concatenate([
-            pose,
-            np.stack([np.stack([np.eye(3)] * pose.shape[0], 0)] * 2, 1)
-        ], 1)
-        shape = [768, 768]
-        render = SMPLRender(cfg.RENDER.SMPL_MODEL_PATH)
-
-        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]
-        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]):
-            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:
-            data = data[None]
-        if isinstance(data, torch.Tensor):
-            data = data.cpu().numpy()
-        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
-
+    np.save(output_joints_path, data)
+    
+    cmd3 = ["xvfb-run", "python", "-m", "render", "--joints_path", output_joints_path, "--method", method, "--output_mp4_path", output_mp4_path, "--smpl_model_path", cfg.RENDER.SMPL_MODEL_PATH]
+    subprocess.run(cmd3)
+    
     return output_mp4_path, video_fname, output_npy_path, feats_fname
 
 

mGPT/render/blender/render.py

@@ -0,0 +1,177 @@
+import math
+import os
+import sys
+import smplx
+# import bpy
+import numpy as np
+
+from .camera import Camera
+from .floor import get_trajectory, plot_floor, show_traj
+from .sampler import get_frameidx
+from .scene import setup_scene  # noqa
+from .tools import delete_objs, load_numpy_vertices_into_blender, style_detect
+from .vertices import prepare_vertices
+from mGPT.utils.joints import smplh_to_mmm_scaling_factor
+
+
+def prune_begin_end(data, perc):
+    to_remove = int(len(data) * perc)
+    if to_remove == 0:
+        return data
+    return data[to_remove:-to_remove]
+
+
+def render_current_frame(path):
+    bpy.context.scene.render.filepath = path
+    bpy.ops.render.render(use_viewport=True, write_still=True)
+
+
+def render(npydata,
+           frames_folder,
+           *,
+           mode,
+           model_path,
+           faces_path,
+           gt=False,
+           exact_frame=None,
+           num=8,
+           downsample=True,
+           canonicalize=True,
+           always_on_floor=False,
+           denoising=True,
+           oldrender=True,
+           res="high",
+           init=True,
+           accelerator='gpu',
+           device=[0]):
+    if init:
+        # Setup the scene (lights / render engine / resolution etc)
+        setup_scene(res=res,
+                    denoising=denoising,
+                    oldrender=oldrender,
+                    accelerator=accelerator,
+                    device=device)
+
+    is_mesh, is_smplx, jointstype = style_detect(npydata)
+    
+    if not is_mesh:
+        npydata = npydata * smplh_to_mmm_scaling_factor
+
+    if is_smplx:
+        smplx_model_male = smplx.create(model_path,
+                                        model_type='smplx',
+                                        gender='male',
+                                        ext='npz',
+                                        num_betas=10,
+                                        flat_hand_mean=True,
+                                        use_pca=False)
+        faces_path = smplx_model_male.faces
+
+    
+
+    # Put everything in this folder
+    if mode == "video":
+        if always_on_floor:
+            frames_folder += "_of"
+        os.makedirs(frames_folder, exist_ok=True)
+        # if it is a mesh, it is already downsampled
+        if downsample and not is_mesh:
+            npydata = npydata[::8]
+    elif mode == "sequence":
+        img_name, ext = os.path.splitext(frames_folder)
+        if always_on_floor:
+            img_name += "_of"
+        img_path = f"{img_name}{ext}"
+
+    elif mode == "frame":
+        img_name, ext = os.path.splitext(frames_folder)
+        if always_on_floor:
+            img_name += "_of"
+        img_path = f"{img_name}_{exact_frame}{ext}"
+
+    # remove X% of begining and end
+    # as it is almost always static
+    # in this part
+    if mode == "sequence":
+        perc = 0.2
+        npydata = prune_begin_end(npydata, perc)
+
+    if is_mesh:
+        from .meshes import Meshes
+        data = Meshes(npydata,
+                      gt=gt,
+                      mode=mode,
+                      faces_path=faces_path,
+                      canonicalize=canonicalize,
+                      always_on_floor=always_on_floor,
+                      is_smplx=is_smplx)
+    else:
+        from .joints import Joints
+        data = Joints(npydata,
+                      gt=gt,
+                      mode=mode,
+                      canonicalize=canonicalize,
+                      always_on_floor=always_on_floor,
+                      jointstype=jointstype)
+
+    # Number of frames possible to render
+    nframes = len(data)
+
+    # Show the trajectory
+    show_traj(data.trajectory)
+
+    # Create a floor
+    plot_floor(data.data, big_plane=False)
+
+    # initialize the camera
+    camera = Camera(first_root=data.get_root(0), mode=mode, is_mesh=is_mesh)
+
+    frameidx = get_frameidx(mode=mode,
+                            nframes=nframes,
+                            exact_frame=exact_frame,
+                            frames_to_keep=num)
+
+    nframes_to_render = len(frameidx)
+
+    # center the camera to the middle
+    if mode == "sequence":
+        camera.update(data.get_mean_root())
+
+    imported_obj_names = []
+    for index, frameidx in enumerate(frameidx):
+        if mode == "sequence":
+            frac = index / (nframes_to_render - 1)
+            mat = data.get_sequence_mat(frac)
+        else:
+            mat = data.mat
+            camera.update(data.get_root(frameidx))
+
+        islast = index == (nframes_to_render - 1)
+
+        objname = data.load_in_blender(frameidx, mat)
+        name = f"{str(index).zfill(4)}"
+
+        if mode == "video":
+            path = os.path.join(frames_folder, f"frame_{name}.png")
+        else:
+            path = img_path
+
+        if mode == "sequence":
+            imported_obj_names.extend(objname)
+        elif mode == "frame":
+            camera.update(data.get_root(frameidx))
+
+        if mode != "sequence" or islast:
+            render_current_frame(path)
+            delete_objs(objname)
+
+    bpy.ops.wm.save_as_mainfile(filepath=frames_folder.replace('.png','.blend').replace('_frames','.blend'))
+
+    # remove every object created
+    delete_objs(imported_obj_names)
+    delete_objs(["Plane", "myCurve", "Cylinder"])
+
+    if mode == "video":
+        return frames_folder
+    else:
+        return img_path

render.py

@@ -0,0 +1,76 @@
+
+from argparse import ArgumentParser
+import numpy as np
+import OpenGL.GL as gl
+import imageio
+import cv2
+import random
+import torch
+import moviepy.editor as mp
+from scipy.spatial.transform import Rotation as RRR
+import mGPT.render.matplot.plot_3d_global as plot_3d
+from mGPT.render.pyrender.hybrik_loc2rot import HybrIKJointsToRotmat
+from mGPT.render.pyrender.smpl_render import SMPLRender
+
+
+if __name__ == '__main__':
+
+    parser = ArgumentParser()
+    parser.add_argument('--joints_path', type=str, help='Path to joints data')
+    parser.add_argument('--method', type=str, help='Method for rendering')
+    parser.add_argument('--output_mp4_path', type=str, help='Path to output MP4 file')
+    parser.add_argument('--smpl_model_path', type=str, help='Path to SMPL model')
+
+    args = parser.parse_args()
+
+    joints_path = args.joints_path
+    method = args.method
+    output_mp4_path = args.output_mp4_path
+    smpl_model_path = args.smpl_model_path
+
+
+    data = np.load(joints_path)
+
+    if method == 'slow':
+        if len(data.shape) == 4:
+            data = data[0]
+        data = data - data[0, 0]
+        pose_generator = HybrIKJointsToRotmat()
+        pose = pose_generator(data)
+        pose = np.concatenate([
+            pose,
+            np.stack([np.stack([np.eye(3)] * pose.shape[0], 0)] * 2, 1)
+        ], 1)
+        shape = [768, 768]
+        render = SMPLRender(smpl_model_path)
+
+        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]
+        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]):
+            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:
+            data = data[None]
+        if isinstance(data, torch.Tensor):
+            data = data.cpu().numpy()
+        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