First commit

app.py

@@ -0,0 +1,265 @@
+from functools import partial
+import os
+
+import torch
+import numpy as np
+import gradio as gr
+import gdown
+
+from load import load_model, load_json
+from load import load_unit_motion_embs_splits, load_keyids_splits
+
+
+EXAMPLES = [
+    "A person is walking in a circle",
+    "A person is jumping rope",
+    "Someone is doing a backflip",
+    "A person is doing a moonwalk",
+    "A person walks forward and then turns back",
+    "Picking up an object",
+    "A person is swimming in the sea",
+    "A human is squatting",
+    "Someone is jumping with one foot",
+    "A person is chopping vegetables",
+    "Someone walks backward",
+    "Somebody is ascending a staircase",
+    "A person is sitting down",
+    "A person is taking the stairs",
+    "Someone is doing jumping jacks",
+    "The person walked forward and is picking up his toolbox",
+    "The person angrily punching the air."
+]
+
+# Show closest text in the training
+
+
+# css to make videos look nice
+CSS = """
+video {
+    position: relative;
+    margin: 0;
+    box-shadow: var(--block-shadow);
+    border-width: var(--block-border-width);
+    border-color: var(--block-border-color);
+    border-radius: var(--block-radius);
+    background: var(--block-background-fill);
+    width: 100%;
+    line-height: var(--line-sm);
+}
+"""
+
+
+def humanml3d_keyid_to_babel_rendered_url(h3d_index, amass_to_babel, keyid):
+    # Don't show the mirrored version of HumanMl3D
+    if "M" in keyid:
+        return None
+
+    dico = h3d_index[keyid]
+    path = dico["path"]
+
+    # HumanAct12 motions are not rendered online
+    # so we skip them for now
+    if "humanact12" in path:
+        return None
+
+    # This motion is not rendered in BABEL
+    # so we skip them for now
+    if path not in amass_to_babel:
+        return None
+
+    babel_id = amass_to_babel[path].zfill(6)
+    url = f"https://babel-renders.s3.eu-central-1.amazonaws.com/{babel_id}.mp4"
+
+    # For the demo, we retrieve from the first annotation only
+    ann = dico["annotations"][0]
+    start = ann["start"]
+    end = ann["end"]
+    text = ann["text"]
+
+    data = {
+        "url": url,
+        "start": start,
+        "end": end,
+        "text": text,
+        "keyid": keyid,
+        "babel_id": babel_id
+    }
+
+    return data
+
+
+def retrieve(model, keyid_to_url, all_unit_motion_embs, all_keyids, text, splits=["test"], nmax=8):
+    unit_motion_embs = torch.cat([all_unit_motion_embs[s] for s in splits])
+    keyids = np.concatenate([all_keyids[s] for s in splits])
+
+    scores = model.compute_scores(text, unit_embs=unit_motion_embs)
+
+    sorted_idxs = np.argsort(-scores)
+    best_keyids = keyids[sorted_idxs]
+    best_scores = scores[sorted_idxs]
+
+    datas = []
+    for keyid, score in zip(best_keyids, best_scores):
+        if len(datas) == nmax:
+            break
+
+        data = keyid_to_url(keyid)
+        if data is None:
+            continue
+        data["score"] = round(float(score), 2)
+        datas.append(data)
+    return datas
+
+
+# HTML component
+def get_video_html(url, video_id, start=None, end=None, score=None, width=350, height=350):
+    trim = ""
+    if start is not None:
+        if end is not None:
+            trim = f"#t={start},{end}"
+        else:
+            trim = f"#t={start}"
+
+    score_t = ""
+    if score is not None:
+        score_t = f'title="Score = {score}"'
+
+    video_html = f'''
+<video preload="auto" muted playsinline onpause="this.load()"
+autoplay loop disablepictureinpicture id="{video_id}" width="{width}" height="{height}" {score_t}>
+  <source src="{url}{trim}" type="video/mp4">
+  Your browser does not support the video tag.
+</video>
+'''
+    return video_html
+
+
+def retrive_component(retrieve_function, text, splits, nvids, n_component=16):
+    # cannot produce more than n_compoenent
+    nvids = min(nvids, n_component)
+    if not splits:
+        return [None for _ in range(n_component)]
+
+    splits_l = [x.lower() for x in splits]
+    datas = retrieve_function(text, splits=splits_l, nmax=nvids)
+    htmls = [
+        get_video_html(
+            url["url"], idx, start=url["start"],
+            end=url["end"], score=url["score"]
+        )
+        for idx, url in enumerate(datas)
+    ]
+    # get n_component exactly if asked less
+    # pad with dummy blocks
+    htmls = htmls + [None for _ in range(max(0, n_component-nvids))]
+    return htmls
+
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # LOADING
+    model = load_model(device)
+    splits = ["train", "val", "test"]
+    all_unit_motion_embs = load_unit_motion_embs_splits(splits, device)
+    all_keyids = load_keyids_splits(splits)
+
+    h3d_index = load_json("amass-annotations/humanml3d.json")
+    amass_to_babel = load_json("amass-annotations/amass_to_babel.json")
+
+    keyid_to_url = partial(humanml3d_keyid_to_babel_rendered_url, h3d_index, amass_to_babel)
+    retrieve_function = partial(retrieve, model, keyid_to_url, all_unit_motion_embs, all_keyids)
+
+    # DEMO
+    theme = gr.themes.Default(primary_hue="blue", secondary_hue="gray")
+    retrive_and_show = partial(retrive_component, retrieve_function)
+
+    default_text = "A person is "
+
+    with gr.Blocks(css=CSS, theme=theme) as demo:
+        title = "<h1 style='text-align: center'>TMR: Text-to-Motion Retrieval Using Contrastive 3D Human Motion Synthesis </h1>"
+        gr.Markdown(title)
+
+        authors = """
+        <h2 style='text-align: center'>
+        <a href="https://mathis.petrovich.fr" target="_blank"><nobr>Mathis Petrovich</nobr></a> &emsp;
+        <a href="https://ps.is.mpg.de/~black" target="_blank"><nobr>Michael J. Black</nobr></a> &emsp;
+	<a href="https://imagine.enpc.fr/~varolg" target="_blank"><nobr>G&uumll Varol</nobr></a>
+	</h2>
+        """
+        gr.Markdown(authors)
+
+        conf = """
+        <h2 style='text-align: center'>
+	<nobr>arXiv 2023</nobr>
+	</h2>
+        """
+        gr.Markdown(conf)
+
+        videos = []
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Column(scale=2):
+                    text = gr.Textbox(placeholder="Type in natural language, the motion to retrieve",
+                                      show_label=True, label="Text prompt", value=default_text)
+                with gr.Column(scale=1):
+                    btn = gr.Button("Retrieve", variant='primary')
+                    clear = gr.Button("Clear", variant='secondary')
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        splits = gr.Dropdown(["Train", "Val", "Test"],
+                                             value=["Test"], multiselect=True, label="Splits",
+                                             info="HumanML3D data used for the motion database")
+                    with gr.Column(scale=1):
+                        nvideo_slider = gr.Slider(minimum=4, maximum=16, step=4, value=8, label="Number of videos")
+            with gr.Column(scale=2):
+                examples = gr.Examples(examples=EXAMPLES, inputs=text, examples_per_page=15)
+
+        i = -1
+        # should indent
+        for _ in range(4):
+            with gr.Row():
+                for _ in range(4):
+                    i += 1
+                    with gr.Column():
+                        video = gr.HTML()
+                        videos.append(video)
+
+        def check_error(splits):
+            if not splits:
+                raise gr.Error("At least one split should be selected!")
+            return splits
+
+        btn.click(fn=retrive_and_show, inputs=[text, splits, nvideo_slider], outputs=videos).then(
+            fn=check_error, inputs=splits
+        )
+
+        text.submit(fn=retrive_and_show, inputs=[text, splits, nvideo_slider], outputs=videos).then(
+            fn=check_error, inputs=splits
+        )
+
+        def keep_test(splits):
+            if len(splits) == 0:
+                return ["Test"]
+            return splits
+
+        def clear_videos():
+            return [None for x in range(16)] + [default_text]
+
+        clear.click(fn=clear_videos, outputs=videos + [text])
+    demo.launch()
+
+
+def prepare():
+    if not os.path.exists("data"):
+        gdown.download_folder("https://drive.google.com/drive/folders/1MgPFgHZ28AMd01M1tJ7YW_1-ut3-4j08", use_cookies=False)
+
+
+if __name__ == "__main__":
+    prepare()
+    main()
+
+# new
+# A person is walking slowly

load.py

@@ -0,0 +1,53 @@
+import os
+import orjson
+import torch
+import numpy as np
+from model import TMR_textencoder
+
+EMBS = "data/unit_motion_embs"
+
+
+def load_json(path):
+    with open(path, "rb") as ff:
+        return orjson.loads(ff.read())
+
+
+def load_keyids(split):
+    path = os.path.join(EMBS, f"{split}.keyids")
+    with open(path) as ff:
+        keyids = np.array([x.strip() for x in ff.readlines()])
+    return keyids
+
+
+def load_keyids_splits(splits):
+    return {
+        split: load_keyids(split)
+        for split in splits
+    }
+
+
+def load_unit_motion_embs(split, device):
+    path = os.path.join(EMBS, f"{split}_motion_embs_unit.npy")
+    tensor = torch.from_numpy(np.load(path)).to(device)
+    return tensor
+
+
+def load_unit_motion_embs_splits(splits, device):
+    return {
+        split: load_unit_motion_embs(split, device)
+        for split in splits
+    }
+
+
+def load_model(device):
+    text_params = {
+        'latent_dim': 256, 'ff_size': 1024, 'num_layers': 6, 'num_heads': 4,
+        'activation': 'gelu', 'modelpath': 'distilbert-base-uncased'
+    }
+    "unit_motion_embs"
+    model = TMR_textencoder(**text_params)
+    state_dict = torch.load("data/textencoder.pt", map_location=device)
+    # load values for the transformer only
+    model.load_state_dict(state_dict, strict=False)
+    model = model.eval()
+    return model

model.py

@@ -0,0 +1,129 @@
+from typing import List, Union
+import torch.nn as nn
+import os
+
+import torch
+import numpy as np
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+from transformers import logging
+from torch.nn.functional import normalize
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+
+        self.register_buffer('pe', pe, persistent=False)
+
+    def forward(self, x):
+        return x + self.pe[:x.shape[0], :]
+
+
+class TMR_textencoder(nn.Module):
+    def __init__(self, modelpath: str, latent_dim: int, ff_size: int,
+                 num_layers: int, num_heads: int, activation: str, **kwargs) -> None:
+        super().__init__()
+
+        logging.set_verbosity_error()
+
+        # Tokenizer
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        self.tokenizer = AutoTokenizer.from_pretrained(modelpath)
+
+        # Text model
+        self.text_model = AutoModel.from_pretrained(modelpath)
+        # Then configure the model
+        self.text_encoded_dim = self.text_model.config.hidden_size
+
+        # Projection of the text-outputs into the latent space
+        self.projection = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(self.text_encoded_dim, latent_dim)
+        )
+
+        self.mu_token = nn.Parameter(torch.randn(latent_dim))
+        self.logvar_token = nn.Parameter(torch.randn(latent_dim))
+        self.sequence_pos_encoding = PositionalEncoding(latent_dim)
+
+        seq_trans_encoder_layer = nn.TransformerEncoderLayer(d_model=latent_dim,
+                                                             nhead=num_heads,
+                                                             dim_feedforward=ff_size,
+                                                             dropout=0.0,
+                                                             activation=activation)
+        self.seqTransEncoder = nn.TransformerEncoder(
+            seq_trans_encoder_layer,
+            num_layers=num_layers
+        )
+
+    def get_last_hidden_state(self, texts: List[str],
+                              return_mask: bool = False
+                              ) -> Union[Tensor, tuple[Tensor, Tensor]]:
+        encoded_inputs = self.tokenizer(texts, return_tensors="pt", padding=True)
+        output = self.text_model(**encoded_inputs.to(self.text_model.device))
+        if not return_mask:
+            return output.last_hidden_state
+        return output.last_hidden_state, encoded_inputs.attention_mask.to(dtype=bool)
+
+    def forward(self, texts: List[str]) -> Tensor:
+        text_encoded, mask = self.get_last_hidden_state(texts, return_mask=True)
+
+        x = self.projection(text_encoded)
+        bs, nframes, _ = x.shape
+        # bs, nframes, totjoints, nfeats = x.shape
+        # Switch sequence and batch_size because the input of
+        # Pytorch Transformer is [Sequence, Batch size, ...]
+        x = x.permute(1, 0, 2)  # now it is [nframes, bs, latent_dim]
+
+        mu_token = torch.tile(self.mu_token, (bs,)).reshape(bs, -1)
+        logvar_token = torch.tile(self.logvar_token, (bs,)).reshape(bs, -1)
+
+        # adding the distribution tokens for all sequences
+        xseq = torch.cat((mu_token[None], logvar_token[None], x), 0)
+
+        # create a bigger mask, to allow attend to mu and logvar
+        token_mask = torch.ones((bs, 2), dtype=bool, device=x.device)
+        aug_mask = torch.cat((token_mask, mask), 1)
+
+        # add positional encoding
+        xseq = self.sequence_pos_encoding(xseq)
+        final = self.seqTransEncoder(xseq, src_key_padding_mask=~aug_mask)
+
+        # only mu for inference
+        mu = final[0]
+        return mu
+
+    # compute score for retrieval
+    def compute_scores(self, texts, unit_embs=None, embs=None):
+        # not both empty
+        assert not (unit_embs is None and embs is None)
+        # not both filled
+        assert not (unit_embs is not None and embs is not None)
+
+        output_str = False
+        # if one input, squeeze the output
+        if isinstance(texts, str):
+            texts = [texts]
+            output_str = True
+
+        # compute unit_embs from embs if not given
+        if embs is not None:
+            unit_embs = normalize(embs)
+
+        with torch.no_grad():
+            latent_unit_texts = normalize(self(texts))
+            # compute cosine similarity between 0 and 1
+            scores = (unit_embs @ latent_unit_texts.T).T/2 + 0.5
+            scores = scores.cpu().numpy()
+
+        if output_str:
+            scores = scores[0]
+
+        return scores