basic functionalities working

.gitignore

@@ -1,3 +1,5 @@
+.cursorignore
+#*.mp4
 .err
 *.out
 /cluster_scripts

app.py

@@ -5,37 +5,27 @@ import torch
 import random
 import os
 from pathlib import Path 
-from aitviewer.headless import HeadlessRenderer
-from aitviewer.configuration import CONFIG as AITVIEWER_CONFIG
+import smplx
+from body_renderer import get_render
+import joblib
 # import cv2
 # import moderngl
 # ctx = moderngl.create_context(standalone=True)
 # print(ctx)
 access_token_smpl = os.environ.get('HF_SMPL_TOKEN')
+os.environ["PYOPENGL_PLATFORM"] = "egl"
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cuda:0' 🤗
 
-DEFAULT_TEXT = "A person is "
+DEFAULT_TEXT = "do it slower "
 
-from aitviewer.models.smpl import SMPLLayer
 def get_smpl_models():
     REPO_ID = 'athn-nik/smpl_models'
     from huggingface_hub import snapshot_download
     return snapshot_download(repo_id=REPO_ID, allow_patterns="smplh*",
                              token=access_token_smpl)
 
-def get_renderer():
-    from aitviewer.headless import HeadlessRenderer
-    from aitviewer.configuration import CONFIG as AITVIEWER_CONFIG
-    smpl_models_path = str(Path(get_smpl_models()))
-    AITVIEWER_CONFIG.update_conf({'playback_fps': 30,
-                                  'auto_set_floor': True,
-                                  'smplx_models': smpl_models_path,
-                                  'z_up': True})
-    return HeadlessRenderer()
-
-
 WEBSITE = ("""<div class="embed_hidden" style="text-align: center;">
     <h1>MotionFix: Text-Driven 3D Human Motion Editing</h1>
     <h3>
@@ -53,13 +43,19 @@ WEBSITE = ("""<div class="embed_hidden" style="text-align: center;">
     </h3>
 </div>
 <div style="display:flex; gap: 0.3rem; justify-content: center; align-items: center;" align="center">
-<a href='https://arxiv.org/abs/'><img src='https://img.shields.io/badge/Arxiv-2405.20340-A42C25?style=flat&logo=arXiv&logoColor=A42C25'></a> 
-<a href='https://arxiv.org/pdf/'><img src='https://img.shields.io/badge/Paper-PDF-yellow?style=flat&logo=arXiv&logoColor=yellow'></a> 
+<a href='https://arxiv.org/pdf/2408.00712'><img src='https://img.shields.io/badge/Arxiv-2405.20340-A42C25?style=flat&logo=arXiv&logoColor=A42C25'></a> 
 <a href='https://motionfix.is.tue.mpg.de'><img src='https://img.shields.io/badge/Project-Page-%23df5b46?style=flat&logo=Google%20chrome&logoColor=%23df5b46'></a> 
-<a href='https://youtube.com/'><img src='https://img.shields.io/badge/YouTube-red?style=flat&logo=youtube&logoColor=white'></a> 
+<a href='https://www.youtube.com/watch?v=cFa6V6Ua-TY'><img src='https://img.shields.io/badge/YouTube-red?style=flat&logo=youtube&logoColor=white'></a> 
 </div>
 """)
 
+CREDITS=("""<div class="embed_hidden" style="text-align: center;">
+    <h3>
+        The renderer of this demo is adapted from the render of 
+        <a href="https://geometry.stanford.edu/projects/humor/" target="_blank" rel="noopener noreferrer">HuMoR</a>
+        with the help of <a href="https://ps.is.mpg.de/person/trakshit" target="_blank" rel="noopener noreferrer">Tithi Rakshit</a> :)
+    </h3>
+""")
 
 WEB_source =  ("""<div class="embed_hidden" style="text-align: center;">
     <h1>Pick a motion to edit!</h1>
@@ -90,7 +86,7 @@ def greet(n):
 def clear():
     return ""
 
-def show_video(input_text):
+def show_video(input_text, key_to_use):
     from normalization import Normalizer
     normalizer = Normalizer()
     from diffusion import create_diffusion
@@ -98,7 +94,12 @@ def show_video(input_text):
     from tmed_denoiser import TMED_denoiser
     model_ckpt = download_models()
     checkpoint = torch.load(model_ckpt)
-
+    motion_to_edit = download_motion_from_dataset(key_to_use)
+    ds_sample = joblib.load(motion_to_edit)
+    source_motion_norm = ds_sample['source_feats_norm'].to('cuda')
+    seqlen_tgt = ds_sample['target_feats_norm'].shape[0]
+    seqlen_src = ds_sample['source_feats_norm'].shape[0]
+    # import ipdb; ipdb.set_trace()
     checkpoint = {k.replace('denoiser.', ''): v for k, v in checkpoint.items()}
     tmed_denoiser = TMED_denoiser().to('cuda')
     tmed_denoiser.load_state_dict(checkpoint, strict=False)
@@ -112,20 +113,16 @@ def show_video(input_text):
                                          noise_schedule='squaredcos_cap_v2',
                                          predict_xstart=True)
     bsz = 1
-    seqlen_tgt = 180
     no_of_texts = len(texts_cond)
     texts_cond = ['']*no_of_texts + texts_cond
     texts_cond = ['']*no_of_texts + texts_cond
     text_emb, text_mask = text_encoder(texts_cond)
 
-    cond_emb_motion = torch.zeros(seqlen_tgt, bsz, 
-                                  512,
-                                  device='cuda')
-    cond_motion_mask = torch.ones((bsz, seqlen_tgt),
+    cond_emb_motion = source_motion_norm.unsqueeze(0).permute(1, 0, 2)
+    cond_motion_mask = torch.ones((bsz, seqlen_src),
                                   dtype=bool, device='cuda')
     mask_target = torch.ones((bsz, seqlen_tgt),
                              dtype=bool, device='cuda')
-
     diff_out = tmed_denoiser._diffusion_reverse(text_emb.to(cond_emb_motion.device),
                                                 text_mask.to(cond_emb_motion.device),
                                                 cond_emb_motion,
@@ -134,31 +131,68 @@ def show_video(input_text):
                                                 diffusion_process,
                                                 init_vec=None,
                                                 init_from='noise',
-                                                gd_text=4.0,
+                                                gd_text=2.0,
                                                 gd_motion=2.0,
                                                 steps_num=300)
     edited_motion = diffout2motion(diff_out, normalizer).squeeze()
-    from renderer import render_motion, color_map, pack_to_render
+    import ipdb; ipdb.set_trace()
     # aitrenderer = get_renderer()
-    AIT_RENDERER = get_renderer()
-    SMPL_LAYER = SMPLLayer(model_type='smplh', ext='npz', gender='neutral')
-    edited_mot_to_render = pack_to_render(rots=edited_motion[..., 3:],
-                                          trans=edited_motion[..., :3])
+    # SMPL_LAYER = SMPLLayer(model_type='smplh', ext='npz', gender='neutral')
+    # edited_mot_to_render = pack_to_render(rots=edited_motion[..., 3:],
+    #                                       trans=edited_motion[..., :3])
+
+    SMPL_MODELS_PATH = str(Path(get_smpl_models()))
+    body_model=smplx.SMPLHLayer(f"{SMPL_MODELS_PATH}/smplh", 
+                                model_type='smplh',
+                                gender='neutral',ext='npz')
+
+    # run_smpl_fwd_verticesbody_model, body_transl, body_orient, body_pose, 
     import random
     xx = random.randint(1, 1000)
-    fname = render_motion(AIT_RENDERER, [edited_mot_to_render],
-                          f"movie_example--{str(xx)}",
-                          pose_repr='aa',
-                          color=[color_map['generated']],
-                          smpl_layer=SMPL_LAYER)
+    # edited_mot_to_render
+    from body_renderer import get_render
+    from transform3d import transform_body_pose
+    edited_motion_aa = transform_body_pose(edited_motion[:, 3:], 
+                                           '6d->aa')
+    if os.path.exists('./output_movie.mp4'):
+        os.remove('./output_movie.mp4')
+    fname = get_render(body_model, 
+               [edited_motion[..., :3].detach().cpu()],
+               [edited_motion_aa[..., :3].detach().cpu()],
+               [edited_motion_aa[..., 3:].detach().cpu()],
+                output_path='./output_movie.mp4',
+                text='', colors=['sky blue'])
+
+    # import ipdb; ipdb.set_trace()
+
+
+    
+    # fname = render_motion(AIT_RENDERER, [edited_mot_to_render],
+    #                       f"movie_example--{str(xx)}",
+    #                       pose_repr='aa',
+    #                       color=[color_map['generated']],
+    #                       smpl_layer=SMPL_LAYER)
+    print(fname)
+    print(os.path.abspath(fname))
     return fname
 
-from huggingface_hub import hf_hub_download, hf_hub_url, cached_download
+from huggingface_hub import hf_hub_download
 
 def download_models():
     REPO_ID = 'athn-nik/example-model'
     return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
 
+def download_motion_from_dataset(key_to_dl):
+    REPO_ID = 'athn-nik/example-model'
+    from huggingface_hub import snapshot_download
+    keytodl = key_to_dl
+    keytodl = '000008'
+    path_for_ds = snapshot_download(repo_id=REPO_ID, 
+                                           allow_patterns=f"dataset_inputs/{keytodl}",
+                                           token=access_token_smpl)
+    path_for_ds_sample = path_for_ds + f'/dataset_inputs/{keytodl}.pth.tar'
+    return path_for_ds_sample
+
 def download_tmr():
     REPO_ID = 'athn-nik/example-model'
     # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
@@ -202,7 +236,7 @@ def random_source_motion(set_to_pick):
     random_key = random.choice(list(current_set.keys()))
     curvid = current_set[random_key]['motion_a']
     text_annot = current_set[random_key]['annotation']
-    return curvid, text_annot
+    return curvid, text_annot, random_key, text_annot
 
 def retrieve_video(retrieve_text):
     tmr_text_encoder = get_tmr_model(download_tmr())
@@ -227,6 +261,7 @@ def retrieve_video(retrieve_text):
     text_annot = all_mots[top_mot]['annotation']
     return curvid, text_annot
 
+
 with gr.Blocks(css="""
     .gradio-row {
         display: flex;
@@ -262,20 +297,21 @@ with gr.Blocks(css="""
     }
 """) as demo:
     gr.Markdown(WEBSITE)
-    
+    random_key_state = gr.State()
+
     with gr.Row(elem_id="gradio-row"):
-        with gr.Column(scale=7, elem_id="gradio-column"):
+        with gr.Column(scale=5, elem_id="gradio-column"):
             gr.Markdown(WEB_source)
             with gr.Row(elem_id="gradio-button-row"):
-                iterative_button = gr.Button("Iterative")
-                retrieve_button = gr.Button("TMRetrieve")
+                # iterative_button = gr.Button("Iterative")
+                # retrieve_button = gr.Button("TMRetrieve")
                 random_button = gr.Button("Random")
 
             with gr.Row(elem_id="gradio-textbox-row"):
                 with gr.Column(scale=5, elem_id="gradio-textbox-with-button"):
-                    retrieve_text = gr.Textbox(placeholder="Type the text for the motion you want to Retrieve:",
-                                               show_label=True, label="Retrieval Text", 
-                                               value=DEFAULT_TEXT)
+                #     retrieve_text = gr.Textbox(placeholder="Type the text for the motion you want to Retrieve:",
+                #                                show_label=True, label="Retrieval Text", 
+                #                                value=DEFAULT_TEXT)
                     clear_button_retrieval = gr.Button("Clear", scale=0)
 
             with gr.Row(elem_id="gradio-textbox-row"):
@@ -288,10 +324,11 @@ with gr.Blocks(css="""
                                    value='all', 
                                    label="Set to pick from", 
                                    info="Motion will be picked from whole dataset or test or train data.")
-
+            # import ipdb; ipdb.set_trace()
             retrieved_video_output = gr.Video(label="Retrieved Motion",
-                                              value=xxx,
+                                            #   value=xxx,
                                               height=360, width=480)
+            
 
         with gr.Column(scale=5, elem_id="gradio-column"):
             gr.Markdown(WEB_target)
@@ -307,8 +344,8 @@ with gr.Blocks(css="""
             video_output = gr.Video(label="Generated Video", height=360, 
                                     width=480)
 
-    def process_and_show_video(input_text):
-        fname = show_video(input_text)
+    def process_and_show_video(input_text, random_key_state):
+        fname = show_video(input_text, random_key_state)
         return fname
 
     def process_and_retrieve_video(input_text):
@@ -318,10 +355,15 @@ with gr.Blocks(css="""
     from retrieval_loader import get_tmr_model
     from dataset_utils import load_motionfix
         
-    edit_button.click(process_and_show_video, inputs=input_text, outputs=video_output)
-    retrieve_button.click(process_and_retrieve_video, inputs=retrieve_text, outputs=[retrieved_video_output, suggested_edit_text])
-    random_button.click(random_source_motion, inputs=set_to_pick, outputs=[retrieved_video_output, suggested_edit_text])
+    edit_button.click(process_and_show_video, inputs=[input_text, random_key_state], outputs=video_output)
+    # retrieve_button.click(process_and_retrieve_video, inputs=retrieve_text, outputs=[retrieved_video_output, suggested_edit_text])
+    random_button.click(random_source_motion, inputs=set_to_pick, 
+                        outputs=[retrieved_video_output,
+                                 suggested_edit_text,
+                                 random_key_state,
+                                 input_text])
     clear_button_edit.click(clear, outputs=input_text)
-    clear_button_retrieval.click(clear, outputs=retrieve_text)
+    # clear_button_retrieval.click(clear, outputs=retrieve_text)
+    gr.Markdown(CREDITS)
 
-demo.launch()
+demo.launch(share=True)

body_renderer.py

@@ -0,0 +1,40 @@
+def get_render(body_model_loaded, 
+               body_trans, 
+               body_orient, body_pose, 
+               output_path, text='', 
+               colors=[]):
+    from renderer.utils import run_smpl_fwd_vertices
+
+    vertices_list=[]
+    if not isinstance(body_trans, list):
+        body_trans = [body_trans]
+    if not isinstance(body_orient, list):
+        body_orient = [body_orient]
+    if not isinstance(body_pose, list):
+        body_pose = [body_pose]
+
+    for trans, orient,pose in zip(body_trans,body_orient,body_pose):
+    
+        vertices= run_smpl_fwd_vertices(body_model_loaded,
+                                        trans,
+                                        orient,
+                                        pose)
+
+        vertices=vertices.vertices
+        vertices = vertices.detach().cpu().numpy()
+        vertices_list.append(vertices)
+
+    #Initialising the renderer
+    from renderer.humor import HumorRenderer
+    fps = 30.0
+    imw = 720 # 480
+    imh = 720 # 360
+    renderer = HumorRenderer(fps=fps, imw=imw, imh=imh)   
+    
+    if len(vertices_list)==2:
+        renderer(vertices_list, output_path, render_pair=True, 
+                 fps=fps,colors=colors)
+    else:
+        renderer(vertices_list[0], output_path, render_pair=False, 
+                 fps=fps,colors=colors)
+    return output_path

renderer.py

@@ -1,138 +0,0 @@
-import os
-import torch
-from transform3d import transform_body_pose
-from aitviewer.headless import HeadlessRenderer
-from gen_utils import rgb, rgba
-
-color_map = {
-             'source_motion': rgba('darkred'),
-             'source': rgba('darkred'),
-             'target_motion': rgba('olivedrab'),
-             'input': rgba('olivedrab'),
-             'target': rgba('olivedrab'),
-             'generation': rgba('purple'),
-             'generated': rgba('steelblue'),
-             'denoised': rgba('purple'),
-             'noised': rgba('darkgrey'),
-             }
-
-
-def pack_to_render(rots, trans, pose_repr='6d'):
-    # make axis-angle
-    # global_orient = transform_body_pose(rots, f"{pose_repr}->aa")
-
-    if rots.is_cuda:
-        rots = rots.detach().cpu()
-    if trans.is_cuda:
-        trans = trans.detach().cpu()
-
-    if pose_repr != 'aa':
-        body_pose = transform_body_pose(rots, f"{pose_repr}->aa")
-    else:
-        body_pose = rots
-    if trans is None:
-        trans = torch.zeros((rots.shape[0], rots.shape[1], 3),
-                             device=rots.device)
-    render_d = {'body_transl': trans,
-                'body_orient': body_pose[..., :3],
-                'body_pose': body_pose[..., 3:]}
-    return render_d
-
-
-def render_motion(renderer: HeadlessRenderer, datum: dict, 
-                  filename: str, pose_repr='6d',
-                  color=(160 / 255, 160 / 255, 160 / 255, 1.0),
-                  return_verts=False, smpl_layer=None) -> None:
-    """
-    Function to render a video of a motion sequence
-    renderer: aitviewer renderer
-    datum: dictionary containing sequence of poses, body translations and body orientations
-        data could be numpy or pytorch tensors
-    filename: the absolute path you want the video to be saved at
-
-    """
-    from aitviewer.headless import HeadlessRenderer
-    from aitviewer.renderables.smpl import SMPLSequence
-
-    if isinstance(datum, dict): datum = [datum]
-    if not isinstance(color, list): 
-        colors = [color] 
-    else:
-        colors = color
-    # assert {'body_transl', 'body_orient', 'body_pose'}.issubset(set(datum[0].keys()))
-    # os.environ['DISPLAY'] = ":11"
-    gender = 'neutral'
-    only_skel = False
-    import sys
-    seqs_of_human_motions = []
-    if smpl_layer is None:
-        from aitviewer.models.smpl import SMPLLayer
-        smpl_layer = SMPLLayer(model_type='smplh', 
-                                ext='npz',
-                                gender=gender)
-    
-    for iid, mesh_seq in enumerate(datum):
-
-        if pose_repr != 'aa':
-            global_orient = transform_body_pose(mesh_seq['body_orient'],
-                                                f"{pose_repr}->aa")
-            body_pose = transform_body_pose(mesh_seq['body_pose'],
-                                            f"{pose_repr}->aa")
-        else:
-            global_orient = mesh_seq['body_orient']
-            body_pose = mesh_seq['body_pose']
-
-        body_transl = mesh_seq['body_transl']
-        sys.stdout.flush()
-
-        old = os.dup(1)
-        os.close(1)
-        os.open(os.devnull, os.O_WRONLY)
-        print(body_pose.shape)
-        print('\n')
-        smpl_template = SMPLSequence(body_pose,
-                                     smpl_layer,
-                                     poses_root=global_orient,
-                                     trans=body_transl,
-                                     color=colors[iid],
-                                     z_up=True)
-        if only_skel:
-            smpl_template.remove(smpl_template.mesh_seq)
-
-        seqs_of_human_motions.append(smpl_template)
-        renderer.scene.add(smpl_template)
-    # camera follows smpl sequence
-    # FIX CAMERA
-    from transform3d import get_z_rot
-    R_z = get_z_rot(global_orient[0], in_format='aa')
-    heading = -R_z[:, 1]
-    xy_facing = body_transl[0] + heading*2.5
-    camera = renderer.lock_to_node(seqs_of_human_motions[0],
-                                    (xy_facing[0], xy_facing[1], 1.5), smooth_sigma=5.0)
-
-    # /FIX CAMERA
-    if len(mesh_seq['body_pose']) == 1:
-        renderer.save_frame(file_path=str(filename) + '.png')
-        sfx = 'png'
-    else:
-        renderer.save_video(video_dir=str(filename), output_fps=30)
-        sfx = 'mp4'
-
-    # aitviewer adds a counter to the filename, we remove it
-    # filename.split('_')[-1].replace('.mp4', '')
-    # os.rename(filename + '_0.mp4', filename[:-4] + '.mp4')
-    if sfx == 'mp4':
-        os.rename(str(filename) + f'_0.{sfx}', str(filename) + f'.{sfx}')
-
-    # empty scene for the next rendering
-    for mesh in seqs_of_human_motions:
-        renderer.scene.remove(mesh)
-    renderer.scene.remove(camera)
-
-    sys.stdout.flush()
-    os.close(1)
-    os.dup(old)
-    os.close(old)
-    renderer.reset()
-    fname = f'{filename}.{sfx}'
-    return fname

renderer/humor.py

@@ -0,0 +1,161 @@
+import numpy as np
+import torch
+from .humor_render_tools.tools import viz_smpl_seq
+from smplx.utils import Struct
+from .video import Video
+import os
+from multiprocessing import Pool
+from tqdm import tqdm
+from multiprocessing import Process
+
+THIS_FOLDER = os.path.dirname(os.path.abspath(__file__))
+FACE_PATH = os.path.join(THIS_FOLDER, "humor_render_tools/smplh.faces")
+FACES = torch.from_numpy(np.int32(np.load(FACE_PATH)))
+
+
+class HumorRenderer:
+    def __init__(self, fps=20.0, **kwargs):
+        self.kwargs = kwargs
+        self.fps = fps
+
+    def __call__(self, vertices, output, render_pair=False,text=None,colors=[],**kwargs):
+        params = self.kwargs | kwargs
+        fps = self.fps
+        if "fps" in params:
+            fps = params.pop("fps")
+
+        if render_pair:
+            render_overlaid(vertices, output,fps,colors,**params)
+        else:
+            render(vertices, output, fps,colors,**params)
+
+        fname = f'{output}.mp4'
+        return fname
+
+
+def render_overlaid(vertices, out_path, fps,colors=[],progress_bar=tqdm,**kwargs):
+    assert isinstance(vertices, list) and len(vertices) == 2
+    # Put the vertices at the floor level
+    # F X N x 3 ===> [F X N x 3, F X N x 3] 
+    ground = vertices[0][..., 2].min()
+    vertices[0][..., 2] -= ground 
+    vertices[1][..., 2] -= ground
+    verts0 = vertices[0]
+    verts1 = vertices[1]
+    
+    import pyrender
+
+    # remove title if it exists
+    kwargs.pop("title", None)
+
+    # vertices: SMPL-H vertices
+    # verts = np.load("interval_2_verts.npy")
+    out_folder = os.path.splitext(out_path)[0]
+
+    verts0 = torch.from_numpy(verts0)
+    body_pred0 = Struct(v=verts0, f=FACES)
+    verts1 = torch.from_numpy(verts1)
+    body_pred1 = Struct(v=verts1, f=FACES)
+
+
+    # out_folder, body_pred, start, end, fps, kwargs = args
+    viz_smpl_seq(
+        pyrender, out_folder, [body_pred0, body_pred1], fps=fps,progress_bar=progress_bar,vertex_color_list=colors, **kwargs
+    )
+
+    video = Video(out_folder, fps=fps)
+    video.save(out_path)
+
+def render(vertices, out_path, fps,colors=[], progress_bar=tqdm,**kwargs):
+    # Put the vertices at the floor level
+    ground = vertices[..., 2].min()
+    vertices[..., 2] -= ground
+
+    import pyrender
+
+    # remove title if it exists
+    kwargs.pop("title", None)
+
+    # vertices: SMPL-H vertices
+    # verts = np.load("interval_2_verts.npy")
+    out_folder = os.path.splitext(out_path)[0]
+
+    verts = torch.from_numpy(vertices)
+    body_pred = Struct(v=verts, f=FACES)
+
+    # out_folder, body_pred, start, end, fps, kwargs = args
+    viz_smpl_seq(
+        pyrender, out_folder, body_pred, fps=fps,progress_bar=progress_bar,vertex_color=colors, **kwargs
+    )
+
+    video = Video(out_folder, fps=fps)
+    video.save(out_path)
+    import shutil
+    shutil.rmtree(out_folder)
+
+def render_offset(args):
+    import pyrender
+
+    out_folder, body_pred, start, end, fps, kwargs = args
+    viz_smpl_seq(
+        pyrender, out_folder, body_pred, start=start, end=end, fps=fps, **kwargs
+    )
+    return 0
+
+
+def render_multiprocess(vertices, out_path, fps, **kwargs):
+    # WIP: does not work yet
+    import ipdb
+
+    ipdb.set_trace()
+    # remove title if it exists
+    kwargs.pop("title", None)
+
+    # vertices: SMPL-H vertices
+    # verts = np.load("interval_2_verts.npy")
+    out_folder = os.path.splitext(out_path)[0]
+
+    verts = torch.from_numpy(vertices)
+    body_pred = Struct(v=verts, f=FACES)
+
+    # faster rendering
+    # by rendering part of the sequence in parallel
+    # still work in progress, use one process for now
+    n_processes = 1
+
+    verts_lst = np.array_split(verts, n_processes)
+    len_split = [len(x) for x in verts_lst]
+    starts = [0] + np.cumsum([x for x in len_split[:-1]]).tolist()
+    ends = np.cumsum([x for x in len_split]).tolist()
+    out_folders = [out_folder for _ in range(n_processes)]
+    fps_s = [fps for _ in range(n_processes)]
+    kwargs_s = [kwargs for _ in range(n_processes)]
+    body_pred_s = [body_pred for _ in range(n_processes)]
+
+    arguments = [out_folders, body_pred_s, starts, ends, fps_s, kwargs_s]
+    # sanity
+    # lst = [verts[start:end] for start, end in zip(starts, ends)]
+    # assert (torch.cat(lst) == verts).all()
+
+    processes = []
+    for _, args in zip(range(n_processes), zip(*arguments)):
+        process = Process(target=render_offset, args=(args,))
+        process.start()
+        processes.append(process)
+
+    for process in processes:
+        process.join()
+
+    if False:
+        # start 4 worker processes
+        with Pool(processes=n_processes) as pool:
+            # print "[0, 1, 4,..., 81]"
+            # print same numbers in arbitrary order
+            print(f"0/{n_processes} rendered")
+            i = 0
+            for _ in pool.imap_unordered(render_offset, zip(*arguments)):
+                i += 1
+                print(f"i/{n_processes} rendered")
+
+    video = Video(out_folder, fps=fps)
+    video.save(out_path)

renderer/humor_render_tools/mesh_viewer.py

@@ -0,0 +1,837 @@
+# adapted from HuMoR
+import os
+import time
+
+# import math
+from tqdm import tqdm
+import numpy as np
+import trimesh
+
+# import pyrender
+import sys
+import cv2
+
+from .parameters import colors
+
+# import pyglet
+
+__all__ = ["MeshViewer"]
+
+COMPRESS_PARAMS = [cv2.IMWRITE_PNG_COMPRESSION, 9]
+
+
+def pause_play_callback(pyrender_viewer, mesh_viewer):
+    mesh_viewer.is_paused = not mesh_viewer.is_paused
+
+
+def step_callback(pyrender_viewer, mesh_viewer, step_size):
+    mesh_viewer.animation_frame_idx = (
+        mesh_viewer.animation_frame_idx + step_size
+    ) % mesh_viewer.animation_len
+
+
+class MeshViewer(object):
+    def __init__(
+        self,
+        pyrender,
+        width=1200,
+        height=800,
+        use_offscreen=False,
+        follow_camera=False,
+        camera_intrinsics=None,
+        img_extn="png",
+        default_cam_offset=[0.0, 4.0, 1.25],
+        default_cam_rot=None,
+    ):
+        super().__init__()
+
+        self.pyrender = pyrender
+        self.use_offscreen = use_offscreen
+        self.follow_camera = follow_camera
+        # render settings for offscreen
+        self.render_wireframe = False
+        self.render_RGBA = False
+        self.render_path = "./render_out"
+        self.img_extn = img_extn
+
+        # mesh sequences to animate
+        self.animated_seqs = []  # the actual sequence of pyrender meshes
+        self.animated_seqs_type = []
+        self.animated_nodes = []  # the nodes corresponding to each sequence
+        self.light_nodes = []
+        # they must all be the same length (set based on first given sequence)
+        self.animation_len = -1
+        # current index in the animation sequence
+        self.animation_frame_idx = 0
+        # track render time to keep steady framerate
+        self.animation_render_time = time.time()
+        # background image sequence
+        self.img_seq = None
+        self.cur_bg_img = None
+        # person mask sequence
+        self.mask_seq = None
+        self.cur_mask = None
+
+        self.single_frame = False
+
+        self.mat_constructor = self.pyrender.MetallicRoughnessMaterial
+        self.trimesh_to_pymesh = self.pyrender.Mesh.from_trimesh
+
+        self.scene = self.pyrender.Scene(
+            bg_color=colors["white"], ambient_light=(0.3, 0.3, 0.3)
+        )
+
+        self.default_cam_offset = np.array(default_cam_offset)
+        self.default_cam_rot = np.array(default_cam_rot)
+
+        self.default_cam_pose = np.eye(4)
+        if default_cam_rot is None:
+            self.default_cam_pose = trimesh.transformations.rotation_matrix(
+                np.radians(180), (0, 0, 1)
+            )
+            self.default_cam_pose = np.dot(
+                trimesh.transformations.rotation_matrix(np.radians(-90), (1, 0, 0)),
+                self.default_cam_pose,
+            )
+        else:
+            self.default_cam_pose[:3, :3] = self.default_cam_rot
+        self.default_cam_pose[:3, 3] = self.default_cam_offset
+
+        self.use_intrins = False
+        if camera_intrinsics is None:
+            pc = self.pyrender.PerspectiveCamera(
+                yfov=np.pi / 3.0, aspectRatio=float(width) / height
+            )
+            camera_pose = self.get_init_cam_pose()
+            self.camera_node = self.scene.add(pc, pose=camera_pose, name="pc-camera")
+
+            light = self.pyrender.DirectionalLight(color=np.ones(3), intensity=1.0)
+            self.scene.add(light, pose=self.default_cam_pose)
+        else:
+            self.use_intrins = True
+            fx, fy, cx, cy = camera_intrinsics
+            camera_pose = np.eye(4)
+            camera_pose = np.array([1.0, -1.0, -1.0, 1.0]).reshape(-1, 1) * camera_pose
+            camera = self.pyrender.camera.IntrinsicsCamera(fx=fx, fy=fy, cx=cx, cy=cy)
+            self.camera_node = self.scene.add(
+                camera, pose=camera_pose, name="pc-camera"
+            )
+
+            light = self.pyrender.DirectionalLight(color=np.ones(3), intensity=1.0)
+            self.scene.add(light, pose=camera_pose)
+
+            self.set_background_color([1.0, 1.0, 1.0, 0.0])
+
+        self.figsize = (width, height)
+
+        # key callbacks
+        self.is_paused = False
+        registered_keys = dict()
+        registered_keys["p"] = (pause_play_callback, [self])
+        registered_keys["."] = (step_callback, [self, 1])
+        registered_keys[","] = (step_callback, [self, -1])
+
+        if self.use_offscreen:
+            self.viewer = self.pyrender.OffscreenRenderer(
+                *self.figsize, point_size=2.75
+            )
+            self.use_raymond_lighting(3.5)
+        else:
+            self.viewer = self.pyrender.Viewer(
+                self.scene,
+                use_raymond_lighting=(not camera_intrinsics),
+                viewport_size=self.figsize,
+                cull_faces=False,
+                run_in_thread=True,
+                registered_keys=registered_keys,
+            )
+
+    def get_init_cam_pose(self):
+        camera_pose = self.default_cam_pose.copy()
+        return camera_pose
+
+    def set_background_color(self, color=colors["white"]):
+        self.scene.bg_color = color
+
+    def update_camera_pose(self, camera_pose):
+        self.scene.set_pose(self.camera_node, pose=camera_pose)
+
+    def close_viewer(self):
+        if self.viewer.is_active:
+            self.viewer.close_external()
+
+    def set_meshes(self, meshes, group_name="static"):
+        for node in self.scene.get_nodes():
+            if node.name is not None and "%s-mesh" % group_name in node.name:
+                self.scene.remove_node(node)
+
+        for mid, mesh in enumerate(meshes):
+            if isinstance(mesh, trimesh.Trimesh):
+                mesh = self.pyrender.Mesh.from_trimesh(mesh.copy())
+            self.acquire_render_lock()
+            self.scene.add(mesh, "%s-mesh-%2d" % (group_name, mid))
+            self.release_render_lock()
+
+    def set_static_meshes(self, meshes):
+        self.set_meshes(meshes, group_name="static")
+
+    def add_static_meshes(self, meshes):
+        for mid, mesh in enumerate(meshes):
+            if isinstance(mesh, trimesh.Trimesh):
+                mesh = self.pyrender.Mesh.from_trimesh(mesh.copy())
+            self.acquire_render_lock()
+            self.scene.add(mesh, "%s-mesh-%2d" % ("staticadd", mid))
+            self.release_render_lock()
+
+    def add_smpl_vtx_list_seq(
+        self, body_mesh_seq, vtx_list, color=[0.0, 0.0, 1.0], radius=0.015
+    ):
+        vtx_point_seq = []
+        for mesh in body_mesh_seq:
+            vtx_point_seq.append(mesh.vertices[vtx_list])
+        self.add_point_seq(vtx_point_seq, color=color, radius=radius)
+
+    def set_img_seq(self, img_seq):
+        """
+        np array of BG images to be rendered in background.
+        """
+        if not self.use_offscreen:
+            print("Cannot render background image if not rendering offscreen")
+            return
+        # ensure same length as other sequences
+        cur_seq_len = len(img_seq)
+        if self.animation_len != -1:
+            if cur_seq_len != self.animation_len:
+                print(
+                    "Unexpected imgage sequence length, all sequences must be the same length!"
+                )
+                return
+        else:
+            if cur_seq_len > 0:
+                self.animation_len = cur_seq_len
+            else:
+                print("Warning: imge sequence is length 0!")
+                return
+
+        self.img_seq = img_seq
+        # must have alpha to render background
+        self.set_render_settings(RGBA=True)
+
+    def set_mask_seq(self, mask_seq):
+        """
+        np array of masked images to be rendered in background.
+        """
+        if not self.use_offscreen:
+            print("Cannot render background image if not rendering offscreen")
+            return
+        # ensure same length as other sequences
+        cur_seq_len = len(mask_seq)
+        if self.animation_len != -1:
+            if cur_seq_len != self.animation_len:
+                print(
+                    "Unexpected imgage sequence length, all sequences must be the same length!"
+                )
+                return
+        else:
+            if cur_seq_len > 0:
+                self.animation_len = cur_seq_len
+            else:
+                print("Warning: imge sequence is length 0!")
+                return
+
+        self.mask_seq = mask_seq
+
+    def add_point_seq(
+        self,
+        point_seq,
+        color=[1.0, 0.0, 0.0],
+        radius=0.015,
+        contact_seq=None,
+        contact_color=[0.0, 1.0, 0.0],
+        connections=None,
+        connect_color=[0.0, 0.0, 1.0],
+        vel=None,
+        render_static=None,
+    ):
+        """
+        Add a sequence of points that will be visualized as spheres.
+
+        - points : List of Nx3 numpy arrays of point locations to visualize as sequence.
+        - color : list of 3 RGB values
+        - radius : radius of each point
+        - contact_seq : an array of num_frames x num_points indicatin "contacts" i.e. points that should be colored
+                        differently at different time steps.
+        - connections : array of point index pairs, draws a cylinder between each pair to create skeleton
+        - vel : list of Nx3 numpy arrays for the velocities of corresponding sequence points
+        """
+        # ensure same length as other sequences
+        cur_seq_len = len(point_seq)
+        if self.animation_len != -1:
+            if cur_seq_len != self.animation_len:
+                print(
+                    "Unexpected sequence length, all sequences must be the same length!"
+                )
+                return
+        else:
+            if cur_seq_len > 0:
+                self.animation_len = cur_seq_len
+            else:
+                print("Warning: points sequence is length 0!")
+                return
+
+        num_joints = point_seq[0].shape[0]
+        if contact_seq is not None and contact_seq.shape[1] != num_joints:
+            print(num_joints)
+            print(contact_seq.shape)
+            print(
+                "Contact sequence must have the same number of points as the input joints!"
+            )
+            return
+        if contact_seq is not None and contact_seq.shape[0] != cur_seq_len:
+            print(
+                "Contact sequence must have the same number of frames as the input sequence!"
+            )
+            return
+
+        # add skeleton
+        if connections is not None:
+            pyrender_skeleton_seq = []
+            for pid, points in enumerate(point_seq):
+                if pid % 200 == 0:
+                    print(
+                        "Caching pyrender connections mesh %d/%d..."
+                        % (pid, len(point_seq))
+                    )
+
+                cyl_mesh_list = []
+                for point_pair in connections:
+                    # print(point_pair)
+                    p1 = points[point_pair[0]]
+                    p2 = points[point_pair[1]]
+                    if np.linalg.norm(p1 - p2) < 1e-6:
+                        segment = np.array([[-1.0, -1.0, -1.0], [-1.01, -1.01, -1.01]])
+                    else:
+                        segment = np.array([p1, p2])
+                    # print(segment)
+
+                    cyl_mesh = trimesh.creation.cylinder(
+                        radius * 0.35, height=None, segment=segment
+                    )
+                    cyl_mesh.visual.vertex_colors = connect_color
+                    cyl_mesh_list.append(cyl_mesh.copy())
+
+                # combine
+                m = self.pyrender.Mesh.from_trimesh(cyl_mesh_list)
+                pyrender_skeleton_seq.append(m)
+
+            if render_static is None:
+                self.add_pyrender_mesh_seq(pyrender_skeleton_seq)
+            else:
+                self.add_static_meshes(
+                    [
+                        pyrender_skeleton_seq[i]
+                        for i in range(len(pyrender_skeleton_seq))
+                        if i % render_static == 0
+                    ]
+                )
+
+        # add velocities
+        if vel is not None:
+            print("Caching pyrender velocities mesh...")
+            pyrender_vel_seq = []
+
+            point_vel_pairs = zip(point_seq, vel)
+            for pid, point_vel_pair in enumerate(point_vel_pairs):
+                cur_point_seq, cur_vel_seq = point_vel_pair
+
+                cyl_mesh_list = []
+                for cur_point, cur_vel in zip(cur_point_seq, cur_vel_seq):
+                    p1 = cur_point
+                    p2 = cur_point + cur_vel * 0.1
+                    segment = np.array([p1, p2])
+                    if np.linalg.norm(p1 - p2) < 1e-6:
+                        continue
+                    cyl_mesh = trimesh.creation.cylinder(
+                        radius * 0.1, height=None, segment=segment
+                    )
+                    cyl_mesh.visual.vertex_colors = [0.0, 0.0, 1.0]
+                    cyl_mesh_list.append(cyl_mesh.copy())
+
+                # combine
+                m = self.pyrender.Mesh.from_trimesh(cyl_mesh_list)
+                pyrender_vel_seq.append(m)
+
+            if render_static is None:
+                self.add_pyrender_mesh_seq(pyrender_vel_seq)
+            else:
+                self.add_static_meshes(
+                    [
+                        pyrender_vel_seq[i]
+                        for i in range(len(pyrender_vel_seq))
+                        if i % render_static == 0
+                    ]
+                )
+
+        # create spheres with trimesh
+        if contact_seq is None:
+            contact_seq = [
+                np.zeros((point_seq[t].shape[0])) for t in range(cur_seq_len)
+            ]
+        pyrender_non_contact_point_seq = []
+        pyrender_contact_point_seq = []
+        for pid, points in enumerate(point_seq):
+            if pid % 200 == 0:
+                print("Caching pyrender points mesh %d/%d..." % (pid, len(point_seq)))
+
+            # first non-contacting points
+            if len(color) > 3:
+                pyrender_non_contact_point_seq.append(
+                    self.pyrender.Mesh.from_points(points, color[pid])
+                )
+            else:
+                sm = trimesh.creation.uv_sphere(radius=radius)
+                sm.visual.vertex_colors = color
+                non_contact_points = points[contact_seq[pid] == 0]
+                if len(non_contact_points) > 0:
+                    tfs = np.tile(np.eye(4), (len(non_contact_points), 1, 1))
+                    tfs[:, :3, 3] = non_contact_points.copy()
+                    m = self.pyrender.Mesh.from_trimesh(sm.copy(), poses=tfs)
+                    pyrender_non_contact_point_seq.append(m)
+                else:
+                    tfs = np.eye(4).reshape((1, 4, 4))
+                    tfs[0, :3, 3] = np.array([0, 0, 30.0])
+                    pyrender_non_contact_point_seq.append(
+                        self.pyrender.Mesh.from_trimesh(sm.copy(), poses=tfs)
+                    )
+            # then contacting points
+            sm = trimesh.creation.uv_sphere(radius=radius)
+            sm.visual.vertex_colors = contact_color
+            contact_points = points[contact_seq[pid] == 1]
+            if len(contact_points) > 0:
+                tfs = np.tile(np.eye(4), (len(contact_points), 1, 1))
+                tfs[:, :3, 3] = contact_points.copy()
+                m = self.pyrender.Mesh.from_trimesh(sm.copy(), poses=tfs)
+                pyrender_contact_point_seq.append(m)
+            else:
+                tfs = np.eye(4).reshape((1, 4, 4))
+                tfs[0, :3, 3] = np.array([0, 0, 30.0])
+                pyrender_contact_point_seq.append(
+                    self.pyrender.Mesh.from_trimesh(sm.copy(), poses=tfs)
+                )
+
+        if len(pyrender_non_contact_point_seq) > 0:
+            if render_static is None:
+                self.add_pyrender_mesh_seq(
+                    pyrender_non_contact_point_seq, seq_type="point"
+                )
+            else:
+                self.add_static_meshes(
+                    [
+                        pyrender_non_contact_point_seq[i]
+                        for i in range(len(pyrender_non_contact_point_seq))
+                        if i % render_static == 0
+                    ]
+                )
+        if len(pyrender_contact_point_seq) > 0:
+            if render_static is None:
+                self.add_pyrender_mesh_seq(pyrender_contact_point_seq, seq_type="point")
+            else:
+                self.add_static_meshes(
+                    [
+                        pyrender_contact_point_seq[i]
+                        for i in range(len(pyrender_contact_point_seq))
+                        if i % render_static == 0
+                    ]
+                )
+
+    def add_mesh_seq(self, mesh_seq, progress_bar=tqdm):
+        """
+        Add a sequence of trimeshes to render.
+
+        - meshes : List of trimesh.trimesh objects giving each frame of the sequence.
+        """
+
+        # ensure same length as other sequences
+        cur_seq_len = len(mesh_seq)
+        if self.animation_len != -1:
+            if cur_seq_len != self.animation_len:
+                print(
+                    "Unexpected sequence length, all sequences must be the same length!"
+                )
+                return
+        else:
+            if cur_seq_len > 0:
+                self.animation_len = cur_seq_len
+            else:
+                print("Warning: mesh sequence is length 0!")
+                return
+
+        # print("Adding mesh sequence with %d frames..." % (cur_seq_len))
+
+        # create sequence of pyrender meshes and save
+        pyrender_mesh_seq = []
+
+        iterator = enumerate(mesh_seq)
+        if progress_bar is not None:
+            iterator = progress_bar(list(iterator), desc="Import meshes in pyrender")
+
+        for mid, mesh in iterator:
+            if isinstance(mesh, trimesh.Trimesh):
+                mesh = self.pyrender.Mesh.from_trimesh(mesh.copy())
+                pyrender_mesh_seq.append(mesh)
+            else:
+                print("Meshes must be from trimesh!")
+                return
+
+        self.add_pyrender_mesh_seq(pyrender_mesh_seq, seq_type="mesh")
+
+    def add_pyrender_mesh_seq(self, pyrender_mesh_seq, seq_type="default"):
+        # add to the list of sequences to render
+        seq_id = len(self.animated_seqs)
+        self.animated_seqs.append(pyrender_mesh_seq)
+        self.animated_seqs_type.append(seq_type)
+
+        # create the corresponding node in the scene
+        self.acquire_render_lock()
+        anim_node = self.scene.add(pyrender_mesh_seq[0], "anim-mesh-%2d" % (seq_id))
+        self.animated_nodes.append(anim_node)
+        self.release_render_lock()
+
+    def add_ground(
+        self,
+        ground_plane=None,
+        length=25.0,
+        color0=[0.8, 0.9, 0.9],
+        color1=[0.6, 0.7, 0.7],
+        tile_width=0.5,
+        xyz_orig=None,
+        alpha=1.0,
+    ):
+        """
+        If ground_plane is none just places at origin with +z up.
+        If ground_plane is given (a, b, c, d) where a,b,c is the normal, then this is rendered. To more accurately place the floor
+        provid an xyz_orig = [x,y,z] that we expect to be near the point of focus.
+        """
+        color0 = np.array(color0 + [alpha])
+        color1 = np.array(color1 + [alpha])
+        # make checkerboard
+        radius = length / 2.0
+        num_rows = num_cols = int(length / tile_width)
+        vertices = []
+        faces = []
+        face_colors = []
+        for i in range(num_rows):
+            for j in range(num_cols):
+                start_loc = [-radius + j * tile_width, radius - i * tile_width]
+                cur_verts = np.array(
+                    [
+                        [start_loc[0], start_loc[1], 0.0],
+                        [start_loc[0], start_loc[1] - tile_width, 0.0],
+                        [start_loc[0] + tile_width, start_loc[1] - tile_width, 0.0],
+                        [start_loc[0] + tile_width, start_loc[1], 0.0],
+                    ]
+                )
+                cur_faces = np.array([[0, 1, 3], [1, 2, 3]], dtype=int)
+                cur_faces += 4 * (
+                    i * num_cols + j
+                )  # the number of previously added verts
+                use_color0 = (i % 2 == 0 and j % 2 == 0) or (i % 2 == 1 and j % 2 == 1)
+                cur_color = color0 if use_color0 else color1
+                cur_face_colors = np.array([cur_color, cur_color])
+
+                vertices.append(cur_verts)
+                faces.append(cur_faces)
+                face_colors.append(cur_face_colors)
+
+        vertices = np.concatenate(vertices, axis=0)
+        faces = np.concatenate(faces, axis=0)
+        face_colors = np.concatenate(face_colors, axis=0)
+
+        if ground_plane is not None:
+            # compute transform between identity floor and passed in floor
+            a, b, c, d = ground_plane
+            # rotation
+            old_normal = np.array([0.0, 0.0, 1.0])
+            new_normal = np.array([a, b, c])
+            new_normal = new_normal / np.linalg.norm(new_normal)
+            v = np.cross(old_normal, new_normal)
+            ang_sin = np.linalg.norm(v)
+            ang_cos = np.dot(old_normal, new_normal)
+            skew_v = np.array(
+                [[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]]
+            )
+            R = (
+                np.eye(3)
+                + skew_v
+                + np.matmul(skew_v, skew_v) * ((1.0 - ang_cos) / (ang_sin**2))
+            )
+            # translation
+            # project point of focus onto plane
+            if xyz_orig is None:
+                xyz_orig = np.array([0.0, 0.0, 0.0])
+            # project origin onto plane
+            plane_normal = np.array([a, b, c])
+            plane_off = d
+            direction = -plane_normal
+            s = (plane_off - np.dot(plane_normal, xyz_orig)) / np.dot(
+                plane_normal, direction
+            )
+            itsct_pt = xyz_orig + s * direction
+            t = itsct_pt
+
+            # transform floor
+            vertices = np.dot(R, vertices.T).T + t.reshape((1, 3))
+
+        ground_tri = trimesh.creation.Trimesh(
+            vertices=vertices, faces=faces, face_colors=face_colors, process=False
+        )
+        ground_mesh = self.pyrender.Mesh.from_trimesh(ground_tri, smooth=False)
+
+        self.acquire_render_lock()
+        anim_node = self.scene.add(ground_mesh, "ground-mesh")
+        self.release_render_lock()
+
+        # update light nodes (if using raymond lighting) to be in this frame
+        if ground_plane is not None:
+            for lnode in self.light_nodes:
+                new_lpose = np.eye(4)
+                new_lrot = np.dot(R, lnode.matrix[:3, :3])
+                new_ltrans = t
+                new_lpose[:3, :3] = new_lrot
+                new_lpose[:3, 3] = new_ltrans
+                self.acquire_render_lock()
+                self.scene.set_pose(lnode, new_lpose)
+                self.release_render_lock()
+
+    def update_frame(self):
+        """
+        Update frame to show the current self.animation_frame_idx
+        """
+        for seq_idx in range(len(self.animated_seqs)):
+            mesh_mean_list=[]
+            # for 2 meshes this is a list with 2 lists with all the frames inside 
+            # import ipdb; ipdb.set_trace()
+            cur_mesh = self.animated_seqs[seq_idx][self.animation_frame_idx]
+    
+            # render the current frame of eqch sequence
+            self.acquire_render_lock()
+
+            # replace the old mesh
+            anim_node = list(self.scene.get_nodes(name="anim-mesh-%2d" % (seq_idx)))
+            anim_node = anim_node[0]
+            anim_node.mesh = cur_mesh
+            # update camera pc-camera
+            if (
+                self.follow_camera and not self.use_intrins
+            ):  # don't want to reset if we're going from camera view
+                if self.animated_seqs_type[seq_idx] == "mesh":
+                    # import ipdb; ipdb.set_trace()
+                    cam_node = list(self.scene.get_nodes(name="pc-camera"))
+                    cam_node = cam_node[0]
+                    if len(self.animated_seqs) > 1:
+                        mesh_mean_paired=[]
+                        for mesh_seq in self.animated_seqs:
+                            mesh_mean_paired.append(mesh_seq[self.animation_frame_idx])
+                        mesh_mean_paired = [ mesh_obj.primitives[0].positions[np.newaxis] for mesh_obj in mesh_mean_paired]
+                        all_meshes_curr_frame = np.concatenate(mesh_mean_paired)
+                        coord_to_add_camera = np.mean(all_meshes_curr_frame, axis=(0, 1))
+                    else:
+                        coord_to_add_camera = np.mean(cur_mesh.primitives[0].positions, axis=0)
+
+                    camera_pose = self.get_init_cam_pose()
+                    camera_pose[:3, 3] = camera_pose[:3, 3] + np.array(
+                        [coord_to_add_camera[0], coord_to_add_camera[1]+1.5, 0.5]
+                    )
+                    self.scene.set_pose(cam_node, camera_pose)
+
+            self.release_render_lock()
+
+        # update background img
+        if self.img_seq is not None:
+            self.acquire_render_lock()
+            self.cur_bg_img = self.img_seq[self.animation_frame_idx]
+            self.release_render_lock
+
+        # update mask
+        if self.mask_seq is not None:
+            self.acquire_render_lock()
+            self.cur_mask = self.mask_seq[self.animation_frame_idx]
+            self.release_render_lock
+
+    def animate(self, fps=30, start=None, end=None, progress_bar=tqdm):
+        """
+        Starts animating any given mesh sequences. This should be called last after adding
+        all desired components to the scene as it is a blocking operation and will run
+        until the user exits (or the full video is rendered if offline).
+        """
+        if not self.use_offscreen:
+            print("=================================")
+            print("VIEWER CONTROLS")
+            print("p - pause/play")
+            print('"," and "." - step back/forward one frame')
+            print("w - wireframe")
+            print("h - render shadows")
+            print("q - quit")
+            print("=================================")
+
+        # print("Animating...")
+        # frame_dur = 1.0 / float(fps)
+
+        # set up init frame
+        self.update_frame()
+
+        # only support offscreen in this script
+        assert self.use_offscreen
+
+        assert self.animation_frame_idx == 0
+
+        iterator = range(self.animation_len)
+        if progress_bar is not None:
+            iterator = progress_bar(list(iterator), desc="SMPL rendering")
+
+        for frame_idx in iterator:
+            self.animation_frame_idx = frame_idx
+            # render frame
+            if not os.path.exists(self.render_path):
+                os.mkdir(self.render_path)
+                # print("Rendering frames to %s!" % (self.render_path))
+            cur_file_path = os.path.join(
+                self.render_path,
+                "frame_%08d.%s" % (self.animation_frame_idx, self.img_extn),
+            )
+            if start is None or end is None:
+                self.save_snapshot(cur_file_path)
+            else:
+                # only do it for a crop
+                if start <= self.animation_frame_idx < end:
+                    self.save_snapshot(cur_file_path)
+
+            if self.animation_frame_idx + 1 >= self.animation_len:
+                continue  # last iteration anyway
+
+            self.animation_render_time = time.time()
+            # if self.is_paused:
+            # self.update_frame()  # just in case there's a single frame update
+            # continue
+
+            self.animation_frame_idx = self.animation_frame_idx + 1
+            # % self.animation_len
+            self.update_frame()
+
+            # if self.single_frame:
+            # break
+
+        self.animation_frame_idx = 0
+        return True
+
+    def _add_raymond_light(self):
+        DirectionalLight = self.pyrender.light.DirectionalLight
+        Node = self.pyrender.node.Node
+        # from pyrender.light import DirectionalLight
+        # from pyrender.node import Node
+
+        thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0])
+        phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0])
+
+        nodes = []
+
+        for phi, theta in zip(phis, thetas):
+            xp = np.sin(theta) * np.cos(phi)
+            yp = np.sin(theta) * np.sin(phi)
+            zp = np.cos(theta)
+
+            z = np.array([xp, yp, zp])
+            z = z / np.linalg.norm(z)
+            x = np.array([-z[1], z[0], 0.0])
+            if np.linalg.norm(x) == 0:
+                x = np.array([1.0, 0.0, 0.0])
+            x = x / np.linalg.norm(x)
+            y = np.cross(z, x)
+
+            matrix = np.eye(4)
+            matrix[:3, :3] = np.c_[x, y, z]
+            nodes.append(
+                Node(
+                    light=DirectionalLight(color=np.ones(3), intensity=1.0),
+                    matrix=matrix,
+                )
+            )
+        return nodes
+
+    def use_raymond_lighting(self, intensity=1.0):
+        if not self.use_offscreen:
+            sys.stderr.write("Interactive viewer already uses raymond lighting!\n")
+            return
+        for n in self._add_raymond_light():
+            n.light.intensity = intensity / 3.0
+            if not self.scene.has_node(n):
+                self.scene.add_node(n)  # , parent_node=pc)
+
+            self.light_nodes.append(n)
+
+    def set_render_settings(
+        self, wireframe=None, RGBA=None, out_path=None, single_frame=None
+    ):
+        if wireframe is not None and wireframe == True:
+            self.render_wireframe = True
+        if RGBA is not None and RGBA == True:
+            self.render_RGBA = True
+        if out_path is not None:
+            self.render_path = out_path
+        if single_frame is not None:
+            self.single_frame = single_frame
+
+    def render(self):
+        RenderFlags = self.pyrender.constants.RenderFlags
+        # from pyrender.constants import RenderFlags
+
+        flags = RenderFlags.SHADOWS_DIRECTIONAL
+        if self.render_RGBA:
+            flags |= RenderFlags.RGBA
+        if self.render_wireframe:
+            flags |= RenderFlags.ALL_WIREFRAME
+        color_img, depth_img = self.viewer.render(self.scene, flags=flags)
+
+        output_img = color_img
+        if self.cur_bg_img is not None:
+            color_img = color_img.astype(np.float32) / 255.0
+            person_mask = None
+            if self.cur_mask is not None:
+                person_mask = self.cur_mask[:, :, np.newaxis]
+                color_img = color_img * (1.0 - person_mask)
+            valid_mask = (color_img[:, :, -1] > 0)[:, :, np.newaxis]
+            input_img = self.cur_bg_img
+            if color_img.shape[2] == 4:
+                output_img = (
+                    color_img[:, :, :-1] * color_img[:, :, 3:]
+                    + (1.0 - color_img[:, :, 3:]) * input_img
+                )
+            else:
+                output_img = (
+                    color_img[:, :, :-1] * valid_mask + (1 - valid_mask) * input_img
+                )
+
+            output_img = (output_img * 255.0).astype(np.uint8)
+
+        return output_img
+
+    def save_snapshot(self, fname):
+        if not self.use_offscreen:
+            sys.stderr.write(
+                "Currently saving snapshots only works with off-screen renderer!\n"
+            )
+            return
+        color_img = self.render()
+        if color_img.shape[-1] == 4:
+            img_bgr = cv2.cvtColor(color_img, cv2.COLOR_RGBA2BGRA)
+        else:
+            img_bgr = cv2.cvtColor(color_img, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(fname, img_bgr, COMPRESS_PARAMS)
+
+    def acquire_render_lock(self):
+        if not self.use_offscreen:
+            self.viewer.render_lock.acquire()
+
+    def release_render_lock(self):
+        if not self.use_offscreen:
+            self.viewer.render_lock.release()

renderer/humor_render_tools/parameters.py

@@ -0,0 +1,32 @@
+# adapted from HuMoR
+# fmt: off
+smpl_connections = [[11, 8], [8, 5], [5, 2], [2, 0], [10, 7], [7, 4], [4, 1], [1, 0],
+                    [0, 3], [3, 6], [6, 9], [9, 12], [12, 15], [12, 13], [13, 16], [16, 18],
+                    [18, 20], [12, 14], [14, 17], [17, 19], [19, 21]]
+
+VERTEX_COLOR = [0.0390625, 0.4140625, 0.796875]
+
+colors = {
+    'pink': [.7, .7, .9],
+    'purple': [.9, .7, .7],
+    'cyan': [.7, .75, .5],
+    'red': [1.0, 0.0, 0.0],
+    'green': [.0, 1., .0],
+    'yellow': [1., 1., 0],
+    'brown': [.5, .7, .7],
+    'blue': [.0, .0, 1.],
+    'offwhite': [.8, .9, .9],
+    'orange': [.5, .65, .9],
+    'grey': [.7, .7, .7],
+    'black': [0.0, 0.0, 0.0],
+    'white': [1.0, 1.0, 1.0],
+    'yellowg': [0.83, 1.0, 0.0],
+    # 'teal':[0.0, 0.7, 0.6],
+    # 'neon pink': (1.0, 0.07, 0.58),
+    'sky blue': [0.1, 0.5, 0.75],
+    'teal': [0.0, 0.6, 0.5],
+    'neon pink': [0.8, 0.05, 0.45],
+
+    "vertex": VERTEX_COLOR
+}
+# fmt: on

renderer/humor_render_tools/tools.py

@@ -0,0 +1,280 @@
+# adapted from HuMoR
+import torch
+import numpy as np
+from tqdm import tqdm
+import trimesh
+from .parameters import colors, smpl_connections
+from .mesh_viewer import MeshViewer
+
+c2c = lambda tensor: tensor.detach().cpu().numpy()  # noqa
+
+
+def viz_smpl_seq(
+    pyrender,
+    out_path,
+    body,
+    #
+    start=None,
+    end=None,
+    #
+    imw=720,
+    imh=720,
+    fps=20,
+    use_offscreen=True,
+    follow_camera=True,
+    progress_bar=tqdm,
+    #
+    contacts=None,
+    render_body=True,
+    render_joints=False,
+    render_skeleton=False,
+    render_ground=True,
+    ground_plane=None,
+    wireframe=False,
+    RGBA=False,
+    joints_seq=None,
+    joints_vel=None,
+    vtx_list=None,
+    points_seq=None,
+    points_vel=None,
+    static_meshes=None,
+    camera_intrinsics=None,
+    img_seq=None,
+    point_rad=0.015,
+    skel_connections=smpl_connections,
+    img_extn="png",
+    ground_alpha=1.0,
+    body_alpha=None,
+    mask_seq=None,
+    cam_offset=[0.0, 2.2, 0.9],  # [0.0, 4.0, 1.25],
+    ground_color0=[0.8, 0.9, 0.9],
+    ground_color1=[0.6, 0.7, 0.7],
+    # ground_color0=[1.0, 1.0, 1.0],
+    # ground_color1=[0.0, 0.0, 0.0],
+    skel_color=[0.5, 0.5, 0.5],  # [0.0, 0.0, 1.0],
+    joint_rad=0.015,
+    point_color=[0.0, 0.0, 1.0],
+    joint_color=[0.0, 1.0, 0.0],
+    contact_color=[1.0, 0.0, 0.0],
+    vertex_color=[],
+    vertex_color_list = [],
+    # color_vtx0 = colors['neon pink']
+    # color_vtx1 = colors['sky blue']
+    render_bodies_static=None,
+    render_points_static=None,
+    cam_rot=None,
+):
+    """
+    Visualizes the body model output of a smpl sequence.
+    - body : body model output from SMPL forward pass (where the sequence is the batch)
+    - joints_seq : list of torch/numy tensors/arrays
+    - points_seq : list of torch/numpy tensors
+    - camera_intrinsics : (fx, fy, cx, cy)
+    - ground_plane : [a, b, c, d]
+    - render_bodies_static is an integer, if given renders all bodies at once but only every x steps
+    """
+
+    if contacts is not None and torch.is_tensor(contacts):
+        contacts = c2c(contacts)
+        # import ipdb; ipdb.set_trace()
+    if isinstance(body, list):
+        render_pair = True
+    else:
+        render_pair = False
+
+    if render_body or vtx_list is not None:
+        if render_pair:
+            nv = body[0].v.size(1)
+        else:
+            nv = body.v.size(1)
+        vertex_colors = np.tile(vertex_color, (nv, 1))
+    
+        if body_alpha is not None:
+            vtx_alpha = np.ones((vertex_colors.shape[0], 1)) * body_alpha
+            vertex_colors = np.concatenate([vertex_colors, vtx_alpha], axis=1)
+        if render_pair:
+            faces = c2c(body[0].f)
+        else:
+            faces = c2c(body.f)
+        
+        if not vertex_color_list:
+            vertex_color_list= ['neon pink','sky blue']
+            
+        if not vertex_color:
+            vertex_color=colors['sky blue']
+        else:
+            vertex_color=colors[vertex_color[0]]
+            
+        if render_pair:
+
+            color_vtx0=colors[vertex_color_list[0]]
+            color_vtx1=colors[vertex_color_list[1]]
+      
+            
+            body_mesh_seq0 = [
+                trimesh.Trimesh(
+                    vertices=c2c(body[0].v[i]),
+                    faces=faces,
+                    vertex_colors=color_vtx0,
+                    process=False,
+                )
+                for i in range(body[0].v.size(0))
+            ]
+            body_mesh_seq1 = [
+                trimesh.Trimesh(
+                    vertices=c2c(body[1].v[i]),
+                    faces=faces,
+                    vertex_colors=color_vtx1,
+                    process=False,
+                )
+                for i in range(body[1].v.size(0))
+            ]
+        else:
+            body_mesh_seq = [
+                trimesh.Trimesh(
+                    vertices=c2c(body.v[i]),
+                    faces=faces,
+                    vertex_colors=vertex_color,
+                    process=False,
+                )
+                for i in range(body.v.size(0))
+            ]
+
+    if render_joints and joints_seq is None:
+        # only body joints
+        joints_seq = [c2c(body.Jtr[i, :22]) for i in range(body.Jtr.size(0))]
+    elif render_joints and torch.is_tensor(joints_seq[0]):
+        joints_seq = [c2c(joint_frame) for joint_frame in joints_seq]
+
+    if joints_vel is not None and torch.is_tensor(joints_vel[0]):
+        joints_vel = [c2c(joint_frame) for joint_frame in joints_vel]
+    if points_vel is not None and torch.is_tensor(points_vel[0]):
+        points_vel = [c2c(joint_frame) for joint_frame in points_vel]
+    # cam_offset = [0.0, 2.2, 0.2]
+    mv = MeshViewer(
+        pyrender,
+        width=imw,
+        height=imh,
+        use_offscreen=use_offscreen,
+        follow_camera=follow_camera,
+        camera_intrinsics=camera_intrinsics,
+        img_extn=img_extn,
+        default_cam_offset=cam_offset,
+        default_cam_rot=cam_rot,
+    )
+    if render_body and render_bodies_static is None:
+        if render_pair:
+            mv.add_mesh_seq(body_mesh_seq0, progress_bar=progress_bar)
+            mv.add_mesh_seq(body_mesh_seq1, progress_bar=progress_bar)
+        else:
+            mv.add_mesh_seq(body_mesh_seq, progress_bar=progress_bar)
+    elif render_body and render_bodies_static is not None:
+        if render_pair:
+            mv.add_static_meshes(
+                [
+                    body_mesh_seq0[i]
+                    for i in range(len(body_mesh_seq0))
+                    if i % render_bodies_static == 0
+                ]
+            )
+            mv.add_static_meshes(
+                [
+                    body_mesh_seq1[i]
+                    for i in range(len(body_mesh_seq1))
+                    if i % render_bodies_static == 0
+                ]
+            )
+
+        else:
+            mv.add_static_meshes(
+                [
+                    body_mesh_seq[i]
+                    for i in range(len(body_mesh_seq))
+                    if i % render_bodies_static == 0
+                ]
+            )
+    if render_joints and render_skeleton:
+        mv.add_point_seq(
+            joints_seq,
+            color=joint_color,
+            radius=joint_rad,
+            contact_seq=contacts,
+            connections=skel_connections,
+            connect_color=skel_color,
+            vel=joints_vel,
+            contact_color=contact_color,
+            render_static=render_points_static,
+        )
+    elif render_joints:
+        mv.add_point_seq(
+            joints_seq,
+            color=joint_color,
+            radius=joint_rad,
+            contact_seq=contacts,
+            vel=joints_vel,
+            contact_color=contact_color,
+            render_static=render_points_static,
+        )
+
+    if vtx_list is not None:
+    
+        mv.add_smpl_vtx_list_seq(
+            body_mesh_seq, vtx_list, color=[0.0, 0.0, 1.0], radius=0.015
+        )
+
+    if points_seq is not None:
+        if torch.is_tensor(points_seq[0]):
+            points_seq = [c2c(point_frame) for point_frame in points_seq]
+        mv.add_point_seq(
+            points_seq,
+            color=point_color,
+            radius=point_rad,
+            vel=points_vel,
+            render_static=render_points_static,
+        )
+
+    if static_meshes is not None:
+        mv.set_static_meshes(static_meshes)
+
+    if img_seq is not None:
+        mv.set_img_seq(img_seq)
+
+    if mask_seq is not None:
+        mv.set_mask_seq(mask_seq)
+
+    if render_ground:
+        xyz_orig = None
+        if ground_plane is not None:
+            if render_body:
+                if render_pair:
+                    xyz_orig = (body_mesh_seq0[0].vertices[0, :] + body_mesh_seq1[0].vertices[0, :]) / 2
+                else:
+                    xyz_orig = body_mesh_seq[0].vertices[0, :]
+
+            elif render_joints:
+                xyz_orig = joints_seq[0][0, :]
+            elif points_seq is not None:
+                xyz_orig = points_seq[0][0, :]
+
+        mv.add_ground(
+            ground_plane=ground_plane,
+            xyz_orig=xyz_orig,
+            color0=ground_color0,
+            color1=ground_color1,
+            alpha=ground_alpha,
+        )
+
+    mv.set_render_settings(
+        out_path=out_path,
+        wireframe=wireframe,
+        RGBA=RGBA,
+        single_frame=(
+            render_points_static is not None or render_bodies_static is not None
+        ),
+    )  # only does anything for offscreen rendering
+    try:
+        mv.animate(fps=fps, start=start, end=end, progress_bar=progress_bar)
+    except RuntimeError as err:
+        print("Could not render properly with the error: %s" % (str(err)))
+
+    del mv

renderer/matplotlib.py

@@ -0,0 +1,259 @@
+# From TEMOS: temos/render/anim.py
+# Inspired by
+# - https://github.com/anindita127/Complextext2animation/blob/main/src/utils/visualization.py
+# - https://github.com/facebookresearch/QuaterNet/blob/main/common/visualization.py
+
+import os
+import logging
+
+from dataclasses import dataclass
+from typing import List, Tuple, Optional
+import numpy as np
+from src.tools.rifke import canonicalize_rotation
+
+logger = logging.getLogger("matplotlib.animation")
+logger.setLevel(logging.ERROR)
+
+colors = ("black", "magenta", "red", "green", "blue")
+
+KINEMATIC_TREES = {
+    "smpljoints": [
+        [0, 3, 6, 9, 12, 15],
+        [9, 13, 16, 18, 20],
+        [9, 14, 17, 19, 21],
+        [0, 1, 4, 7, 10],
+        [0, 2, 5, 8, 11],
+    ],
+    "guoh3djoints": [  # no hands
+        [0, 3, 6, 9, 12, 15],
+        [9, 13, 16, 18, 20],
+        [9, 14, 17, 19, 21],
+        [0, 1, 4, 7, 10],
+        [0, 2, 5, 8, 11],
+    ],
+}
+
+
+@dataclass
+class MatplotlibRender:
+    jointstype: str = "smpljoints"
+    fps: float = 20.0
+    colors: List[str] = colors
+    figsize: int = 4
+    fontsize: int = 15
+    canonicalize: bool = False
+
+    def __call__(
+        self,
+        joints,
+        output,
+        fps=None,
+        highlights=None,
+        title: str = "",
+        canonicalize=None,
+    ):
+        canonicalize = canonicalize if canonicalize is not None else self.canonicalize
+        fps = fps if fps is not None else self.fps
+        if joints.shape[1] == 24:
+            # remove the hands
+            joints = joints[:, :22]
+
+        render_animation(
+            joints,
+            title=title,
+            highlights=highlights,
+            output=output,
+            jointstype=self.jointstype,
+            fps=self.fps,
+            colors=self.colors,
+            figsize=(self.figsize, self.figsize),
+            fontsize=self.fontsize,
+            canonicalize=canonicalize,
+        )
+
+
+def init_axis(fig, title, radius=1.5):
+    ax = fig.add_subplot(1, 1, 1, projection="3d")
+    ax.view_init(elev=20.0, azim=-60)
+
+    fact = 2
+    ax.set_xlim3d([-radius / fact, radius / fact])
+    ax.set_ylim3d([-radius / fact, radius / fact])
+    ax.set_zlim3d([0, radius])
+
+    ax.set_aspect("auto")
+    ax.set_xticklabels([])
+    ax.set_yticklabels([])
+    ax.set_zticklabels([])
+
+    ax.set_axis_off()
+    ax.grid(b=False)
+
+    ax.set_title(title, loc="center", wrap=True)
+    return ax
+
+
+def plot_floor(ax, minx, maxx, miny, maxy, minz):
+    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+    # Plot a plane XZ
+    verts = [
+        [minx, miny, minz],
+        [minx, maxy, minz],
+        [maxx, maxy, minz],
+        [maxx, miny, minz],
+    ]
+    xz_plane = Poly3DCollection([verts], zorder=1)
+    xz_plane.set_facecolor((0.5, 0.5, 0.5, 1))
+    ax.add_collection3d(xz_plane)
+
+    # Plot a bigger square plane XZ
+    radius = max((maxx - minx), (maxy - miny))
+
+    # center +- radius
+    minx_all = (maxx + minx) / 2 - radius
+    maxx_all = (maxx + minx) / 2 + radius
+
+    miny_all = (maxy + miny) / 2 - radius
+    maxy_all = (maxy + miny) / 2 + radius
+
+    verts = [
+        [minx_all, miny_all, minz],
+        [minx_all, maxy_all, minz],
+        [maxx_all, maxy_all, minz],
+        [maxx_all, miny_all, minz],
+    ]
+    xz_plane = Poly3DCollection([verts], zorder=1)
+    xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+    ax.add_collection3d(xz_plane)
+    return ax
+
+
+def update_camera(ax, root, radius=1.5):
+    fact = 2
+    ax.set_xlim3d([-radius / fact + root[0], radius / fact + root[0]])
+    ax.set_ylim3d([-radius / fact + root[1], radius / fact + root[1]])
+
+
+def render_animation(
+    joints: np.ndarray,
+    output: str = "notebook",
+    highlights: Optional[np.ndarray] = None,
+    jointstype: str = "smpljoints",
+    title: str = "",
+    fps: float = 20.0,
+    colors: List[str] = colors,
+    figsize: Tuple[int] = (4, 4),
+    fontsize: int = 15,
+    canonicalize: bool = False,
+    agg=True,
+):
+    if agg:
+        import matplotlib
+
+        matplotlib.use("Agg")
+
+    if highlights is not None:
+        assert len(highlights) == len(joints)
+
+    assert jointstype in KINEMATIC_TREES
+    kinematic_tree = KINEMATIC_TREES[jointstype]
+
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation
+    import matplotlib.patheffects as pe
+
+    mean_fontsize = fontsize
+
+    # heuristic to change fontsize
+    fontsize = mean_fontsize - (len(title) - 30) / 20
+    plt.rcParams.update({"font.size": fontsize})
+
+    # Z is gravity here
+    x, y, z = 0, 1, 2
+
+    joints = joints.copy()
+
+    if canonicalize:
+        joints = canonicalize_rotation(joints, jointstype=jointstype)
+
+    # Create a figure and initialize 3d plot
+    fig = plt.figure(figsize=figsize)
+    ax = init_axis(fig, title)
+
+    # Create spline line
+    trajectory = joints[:, 0, [x, y]]
+    avg_segment_length = (
+        np.mean(np.linalg.norm(np.diff(trajectory, axis=0), axis=1)) + 1e-3
+    )
+    draw_offset = int(25 / avg_segment_length)
+    (spline_line,) = ax.plot(*trajectory.T, zorder=10, color="white")
+
+    # Create a floor
+    minx, miny, _ = joints.min(axis=(0, 1))
+    maxx, maxy, _ = joints.max(axis=(0, 1))
+    plot_floor(ax, minx, maxx, miny, maxy, 0)
+
+    # Put the character on the floor
+    height_offset = np.min(joints[:, :, z])  # Min height
+    joints = joints.copy()
+    joints[:, :, z] -= height_offset
+
+    # Initialization for redrawing
+    lines = []
+    initialized = False
+
+    def update(frame):
+        nonlocal initialized
+        skeleton = joints[frame]
+
+        root = skeleton[0]
+        update_camera(ax, root)
+
+        hcolors = colors
+        if highlights is not None and highlights[frame]:
+            hcolors = ("red", "red", "red", "red", "red")
+
+        for index, (chain, color) in enumerate(
+            zip(reversed(kinematic_tree), reversed(hcolors))
+        ):
+            if not initialized:
+                lines.append(
+                    ax.plot(
+                        skeleton[chain, x],
+                        skeleton[chain, y],
+                        skeleton[chain, z],
+                        linewidth=6.0,
+                        color=color,
+                        zorder=20,
+                        path_effects=[pe.SimpleLineShadow(), pe.Normal()],
+                    )
+                )
+
+            else:
+                lines[index][0].set_xdata(skeleton[chain, x])
+                lines[index][0].set_ydata(skeleton[chain, y])
+                lines[index][0].set_3d_properties(skeleton[chain, z])
+                lines[index][0].set_color(color)
+
+        left = max(frame - draw_offset, 0)
+        right = min(frame + draw_offset, trajectory.shape[0])
+
+        spline_line.set_xdata(trajectory[left:right, 0])
+        spline_line.set_ydata(trajectory[left:right, 1])
+        spline_line.set_3d_properties(np.zeros_like(trajectory[left:right, 0]))
+        initialized = True
+
+    fig.tight_layout()
+    frames = joints.shape[0]
+    anim = FuncAnimation(fig, update, frames=frames, interval=1000 / fps, repeat=False)
+
+    if output == "notebook":
+        from IPython.display import HTML
+
+        HTML(anim.to_jshtml())
+    else:
+        # anim.save(output, writer='ffmpeg', fps=fps)
+        anim.save(output, fps=fps)
+
+    plt.close()

renderer/utils.py

@@ -0,0 +1,33 @@
+from transform3d import transform_body_pose
+
+def run_smpl_fwd_vertices(body_model, 
+                          body_transl, 
+                          body_orient, body_pose):
+    """
+    Standalone function to run SMPL forward pass.run_smpl_fwd_vertices
+    
+    Parameters:
+    - body_model: The SMPL model instance
+    - body_transl: Translation tensor
+    - body_orient: Orientation tensor
+    - body_pose: Pose tensor
+    - fast: Boolean flag to use fast version
+    
+    Returns:
+    - Vertices from the SMPL forward pass
+    """
+    
+    if len(body_transl.shape) > 2:
+        body_transl = body_transl.flatten(0, 1)
+        body_orient = body_orient.flatten(0, 1)
+        body_pose = body_pose.flatten(0, 1)
+
+    batch_size = body_transl.shape[0]
+    body_model.batch_size = batch_size
+    
+    return body_model(
+        transl=body_transl,
+        body_pose=transform_body_pose(body_pose, 'aa->rot'),
+        global_orient=transform_body_pose(body_orient, 'aa->rot')
+    )
+

renderer/video.py

@@ -0,0 +1,39 @@
+import moviepy.editor as mp
+
+# import moviepy.video.fx.all as vfx
+import os
+
+
+class Video:
+    def __init__(self, frame_path: str, fps: float = 20.0, res="high"):
+        frame_path = str(frame_path)
+        self.fps = fps
+
+        self._conf = {
+            "codec": "libx264",
+            "fps": self.fps,
+            "audio_codec": "aac",
+            "temp_audiofile": "temp-audio.m4a",
+            "remove_temp": True,
+        }
+
+        if res == "low":
+            bitrate = "500k"
+        else:
+            bitrate = "5000k"
+
+        self._conf = {"bitrate": bitrate, "fps": self.fps}
+
+        # Load video
+        # video = mp.VideoFileClip(video1_path, audio=False)
+        # Load with frames
+        frames = [os.path.join(frame_path, x) for x in sorted(os.listdir(frame_path))]
+        video = mp.ImageSequenceClip(frames, fps=fps)
+        self.video = video
+        self.duration = video.duration
+
+    def save(self, out_path):
+        out_path = str(out_path)
+        self.video.subclip(0, self.duration).write_videofile(
+            out_path, verbose=False, logger=None, **self._conf
+        )

requirements.txt

@@ -5,4 +5,5 @@ transformers==4.41.2
 hydra-core
 einops
 roma
-aitviewer
+pyrender
+moviepy