fix first version

app.py

@@ -3,12 +3,15 @@ from geometry_utils import diffout2motion
 import gradio as gr
 import spaces
 import torch
-import random
 import os
 from pathlib import Path 
 import smplx
 from body_renderer import get_render
-import joblib
+import numpy as np
+from download_deps import get_smpl_models, download_models, download_model_config
+from download_deps import download_tmr, download_motionfix, download_motionfix_dataset
+from download_deps import download_embeddings
+from website import CREDITS, WEB_source, WEB_target, WEBSITE
 # import cv2
 # import moderngl
 # ctx = moderngl.create_context(standalone=True)
@@ -21,60 +24,7 @@ print(zero.device) # <-- 'cuda:0' 🤗
 
 DEFAULT_TEXT = "do it slower "
 
-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)
-
-WEBSITE = ("""<div class="embed_hidden" style="text-align: center;">
-    <h1>MotionFix: Text-Driven 3D Human Motion Editing</h1>
-    <h3>
-        <a href="https://is.mpg.de/person/~nathanasiou" target="_blank" rel="noopener noreferrer">Nikos Athanasiou</a><sup>1</sup>,
-        <a href="https://is.mpg.de/person/acseke" target="_blank" rel="noopener noreferrer">Alpar Cseke</a><sup>1</sup>,
-        <br>
-        <a href="https://ps.is.mpg.de/person/mdiomataris" target="_blank" rel="noopener noreferrer">Markos Diomataris</a><sup>1, 3</sup>,
-        <a href="https://is.mpg.de/person/black" target="_blank" rel="noopener noreferrer">Michael J. Black</a><sup>1</sup>,
-        <a href="https://imagine.enpc.fr/~varolg/" target="_blank" rel="noopener noreferrer">G&uuml;l Varol</a><sup>2</sup>
-    </h3>
-    <h3>
-        <sup>1</sup>Max Planck Institute for Intelligent Systems, T&uuml;bingen, Germany;
-        <sup>2</sup>LIGM, &Eacute;cole des Ponts, Univ Gustave Eiffel, CNRS, France,
-        <sup>3</sup>ETH Z&uuml;rich, Switzerland
-    </h3>
-</div>
-<div style="display:flex; gap: 0.3rem; justify-content: center; align-items: center;" align="center">
-<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://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>
-    <h3>
-        Here you should pick a source motion
-        <hr class="double">
-    </h3>
-</div>
-""")
-
-WEB_target =  ("""<div class="embed_hidden" style="text-align: center;">
-    <h1>Now type the text to edit that motion!</h1>
-    <h3>
-        Here you should get the generated motion!
-        <hr class="double">
-    </h3>
-</div>
-""")
+
 @spaces.GPU
 def greet(n):
     print(zero.device) # <-- 'cuda:0' 🤗
@@ -100,16 +50,17 @@ def show_video(input_text, key_to_use):
     # ds_sample = joblib.load(motion_to_edit)
     ds_sample = MFIX_DATASET_DICT[key_to_use]
     from feature_extractor import FEAT_GET_METHODS
-    data_dict_source = {f'{feat}_source': FEAT_GET_METHODS[feat](ds_sample['motion_source'])
+    data_dict_source = {f'{feat}_source': FEAT_GET_METHODS[feat](ds_sample['motion_source'])[None].to('cuda')
                         for feat in infeats}
-    data_dict_target = {f'{feat}_target': FEAT_GET_METHODS[feat](ds_sample['motion_target'])
+    data_dict_target = {f'{feat}_target': FEAT_GET_METHODS[feat](ds_sample['motion_target'])[None].to('cuda')
                         for feat in infeats}
     full_batch = data_dict_source | data_dict_target
-    import ipdb; ipdb.set_trace()
     in_batch = normalizer.norm_and_cat(full_batch, infeats)
-    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]
+    source_motion_norm = in_batch['source']
+    target_motion_norm = in_batch['target']
+
+    seqlen_tgt = source_motion_norm.shape[0]
+    seqlen_src = target_motion_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')
@@ -128,8 +79,8 @@ def show_video(input_text, key_to_use):
     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 = source_motion_norm.unsqueeze(0).permute(1, 0, 2)
+    
+    cond_emb_motion = source_motion_norm
     cond_motion_mask = torch.ones((bsz, seqlen_src),
                                   dtype=bool, device='cuda')
     mask_target = torch.ones((bsz, seqlen_tgt),
@@ -145,7 +96,9 @@ def show_video(input_text, key_to_use):
                                                 gd_text=2.0,
                                                 gd_motion=2.0,
                                                 steps_num=300)
-    edited_motion = diffout2motion(diff_out, normalizer).squeeze()
+    edited_motion = diffout2motion(diff_out.permute(1,0,2), normalizer).squeeze()
+    gt_source = diffout2motion(source_motion_norm.permute(1,0,2),
+                               normalizer).squeeze()
     # import ipdb; ipdb.set_trace()
     # aitrenderer = get_renderer()
     # SMPL_LAYER = SMPLLayer(model_type='smplh', ext='npz', gender='neutral')
@@ -158,26 +111,43 @@ def show_video(input_text, key_to_use):
                                 gender='neutral',ext='npz')
 
     # run_smpl_fwd_verticesbody_model, body_transl, body_orient, body_pose, 
-    import random
-    xx = random.randint(1, 1000)
     # edited_mot_to_render
     from body_renderer import get_render
     from transform3d import transform_body_pose
+    # import ipdb; ipdb.set_trace()
     edited_motion_aa = transform_body_pose(edited_motion[:, 3:], 
                                            '6d->aa')
+    gt_source_aa = transform_body_pose(gt_source[:, 3:], 
+                                           '6d->aa')
+
     if os.path.exists('./output_movie.mp4'):
         os.remove('./output_movie.mp4')
+    from transform3d import rotate_body_degrees
+    gen_motion_trans = edited_motion[..., :3].detach().cpu()
+    gen_motion_rots_aa = edited_motion_aa.detach().cpu()
+    source_motion_trans = gt_source[..., :3].detach().cpu()
+    source_motion_rots_aa = gt_source_aa.detach().cpu()
+    
+    gen_rots_rotated, gen_trans_rotated  = rotate_body_degrees(transform_body_pose(
+                                            gen_motion_rots_aa, 
+                                           'aa->rot'),
+                                             gen_motion_trans, offset=np.pi)
+    src_rots_rotated, src_trans_rotated = rotate_body_degrees(transform_body_pose(
+                                            source_motion_rots_aa, 
+                                           'aa->rot'),
+                                             source_motion_trans, offset=np.pi)
+
+    src_rots_rotated_aa = transform_body_pose(src_rots_rotated, 
+                                             'rot->aa')
+    gen_rots_rotated_aa = transform_body_pose(gen_rots_rotated, 
+                                             'rot->aa')
     fname = get_render(body_model, 
-               [edited_motion[..., :3].detach().cpu()],
-               [edited_motion_aa[..., :3].detach().cpu()],
-               [edited_motion_aa[..., 3:].detach().cpu()],
+               [gen_trans_rotated, src_trans_rotated],
+               [gen_rots_rotated_aa[:, 0], src_rots_rotated_aa[:, 0]],
+               [gen_rots_rotated_aa[:, 1:], src_rots_rotated_aa[:, 1:]],
                 output_path='./output_movie.mp4',
-                text='', colors=['sky blue'])
-
-    # import ipdb; ipdb.set_trace()
-
+                text='', colors=['sky blue', 'red'])
 
-    
     # fname = render_motion(AIT_RENDERER, [edited_mot_to_render],
     #                       f"movie_example--{str(xx)}",
     #                       pose_repr='aa',
@@ -187,57 +157,6 @@ def show_video(input_text, key_to_use):
     print(os.path.abspath(fname))
     return fname
 
-from huggingface_hub import hf_hub_download
-
-def download_models():
-    REPO_ID = 'athn-nik/example-model'
-    return hf_hub_download(REPO_ID, filename="tmed_compressed.ckpt")
-
-def download_model_config():
-    REPO_ID = 'athn-nik/example-model'
-    path_to_config = hf_hub_download(REPO_ID, filename="tmed/.hydra/config.yaml")
-    from omegaconf import OmegaConf
-    model_cfg = OmegaConf.to_yaml(path_to_config)
-    return model_cfg.model.input_feats
-
-
-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")
-    from huggingface_hub import snapshot_download
-    return snapshot_download(repo_id=REPO_ID, allow_patterns="tmr*",
-                             token=access_token_smpl)
-
-def download_motionfix():
-    REPO_ID = 'athn-nik/example-model'
-    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
-    from huggingface_hub import snapshot_download
-    return snapshot_download(repo_id=REPO_ID, allow_patterns="motionfix*",
-                             token=access_token_smpl)
-
-def download_motionfix_dataset():
-    REPO_ID = 'athn-nik/example-model'
-    dataset_downloaded_path = hf_hub_download(REPO_ID, filename="motionfix.pth.tar")
-    dataset_dict = joblib.load(dataset_downloaded_path)
-    return dataset_dict
-
-def download_embeddings():
-    REPO_ID = 'athn-nik/example-model'
-    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
-    from huggingface_hub import snapshot_download
-    return snapshot_download(repo_id=REPO_ID, allow_patterns="embeddings*",
-                             token=access_token_smpl)
 
 
 MFIX_p = download_motionfix() + '/motionfix'

deps/{statistics_bodilex.npy → statistics_motionfix.npy}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a0be87962557d3149203eb4586f3e670c1bd7785765ad8cef9ed91f6277a2c2
+oid sha256:dee69fd95d6bda36218e81cf36671a3e61988a358afce8338e771bcd3f14cef1
 size 4826

download_deps.py

@@ -0,0 +1,60 @@
+import os
+import joblib
+from huggingface_hub import hf_hub_download
+access_token_smpl = os.environ.get('HF_SMPL_TOKEN')
+
+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 download_models():
+    REPO_ID = 'athn-nik/example-model'
+    return hf_hub_download(REPO_ID, filename="tmed_compressed.ckpt")
+
+def download_model_config():
+    REPO_ID = 'athn-nik/example-model'
+    path_to_config = hf_hub_download(REPO_ID, filename="tmed/.hydra/config.yaml")
+    from omegaconf import OmegaConf
+    model_cfg = OmegaConf.load(path_to_config)
+    return model_cfg.model.input_feats
+
+
+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")
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=REPO_ID, allow_patterns="tmr*",
+                             token=access_token_smpl)
+
+def download_motionfix():
+    REPO_ID = 'athn-nik/example-model'
+    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=REPO_ID, allow_patterns="motionfix*",
+                             token=access_token_smpl)
+
+def download_motionfix_dataset():
+    REPO_ID = 'athn-nik/example-model'
+    dataset_downloaded_path = hf_hub_download(REPO_ID, filename="motionfix.pth.tar")
+    dataset_dict = joblib.load(dataset_downloaded_path)
+    return dataset_dict
+
+def download_embeddings():
+    REPO_ID = 'athn-nik/example-model'
+    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=REPO_ID, allow_patterns="embeddings*",
+                             token=access_token_smpl)

normalization.py

@@ -7,7 +7,7 @@ import torch
 import numpy as np
 
 class Normalizer:
-    def __init__(self, statistics_path: str='deps/statistics_bodilex.npy', nfeats: int=207, 
+    def __init__(self, statistics_path: str='deps/statistics_motionfix.npy', nfeats: int=207, 
                  input_feats: List[str] = ["body_transl_delta_pelv",
                                            "body_orient_xy",
                                            "z_orient_delta", "body_pose",

renderer/humor_render_tools/mesh_viewer.py

@@ -640,7 +640,7 @@ class MeshViewer(object):
 
                     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]
+                        [coord_to_add_camera[0], coord_to_add_camera[1] + 1.5, 0.5]
                     )
                     self.scene.set_pose(cam_node, camera_pose)
 

transform3d.py

@@ -48,12 +48,11 @@ def rotate_trans(trans, rotZ, inverse=False):
     return trans_local
 
 
-def rotate_body_canonic(rots, trans, offset=0.0):
+def rotate_body_degrees(rots, trans, offset=0.0):
 
-    '''
+    """
         Rotate the whole body
-    '''
-
+    """
     # rots, trans = data.rots.clone(), data.trans.clone()
     global_poses = rots[..., 0, :, :]
     global_euler = matrix_to_euler_angles(global_poses, "ZYX")
@@ -80,11 +79,6 @@ def rotate_body_canonic(rots, trans, offset=0.0):
     # return RotTransDatastruct(rots=rots, trans=trans)
     return rots, trans
 
-
-
-
-
-
 """
 The transformation matrices returned from the functions in this file assume
 the points on which the transformation will be applied are column vectors.

website.py

@@ -0,0 +1,49 @@
+
+WEBSITE = ("""<div class="embed_hidden" style="text-align: center;">
+    <h1>MotionFix: Text-Driven 3D Human Motion Editing</h1>
+    <h3>
+        <a href="https://is.mpg.de/person/~nathanasiou" target="_blank" rel="noopener noreferrer">Nikos Athanasiou</a><sup>1</sup>,
+        <a href="https://is.mpg.de/person/acseke" target="_blank" rel="noopener noreferrer">Alpar Cseke</a><sup>1</sup>,
+        <br>
+        <a href="https://ps.is.mpg.de/person/mdiomataris" target="_blank" rel="noopener noreferrer">Markos Diomataris</a><sup>1, 3</sup>,
+        <a href="https://is.mpg.de/person/black" target="_blank" rel="noopener noreferrer">Michael J. Black</a><sup>1</sup>,
+        <a href="https://imagine.enpc.fr/~varolg/" target="_blank" rel="noopener noreferrer">G&uuml;l Varol</a><sup>2</sup>
+    </h3>
+    <h3>
+        <sup>1</sup>Max Planck Institute for Intelligent Systems, T&uuml;bingen, Germany;
+        <sup>2</sup>LIGM, &Eacute;cole des Ponts, Univ Gustave Eiffel, CNRS, France,
+        <sup>3</sup>ETH Z&uuml;rich, Switzerland
+    </h3>
+</div>
+<div style="display:flex; gap: 0.3rem; justify-content: center; align-items: center;" align="center">
+<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://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>
+    <h3>
+        Here you should pick a source motion
+        <hr class="double">
+    </h3>
+</div>
+""")
+
+WEB_target =  ("""<div class="embed_hidden" style="text-align: center;">
+    <h1>Now type the text to edit that motion!</h1>
+    <h3>
+        Here you should get the generated motion!
+        <hr class="double">
+    </h3>
+</div>
+""")