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 WEBSITE = """

TMR: Text-to-Motion Retrieval Using Contrastive 3D Human Motion Synthesis

Mathis PetrovichMichael J. BlackGül Varol

arXiv 2023

Description

This space illustrates TMR, a method for text-to-motion retrieval. Given a gallery of 3D human motions (which can be unseen during training) and a text query, the goal is to search for motions which are close to the text query.

""" EXAMPLES = [ "A person is walking slowly", "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 # var(--block-border-color); CSS = """ .retrieved_video { position: relative; margin: 0; box-shadow: var(--block-shadow); border-width: var(--block-border-width); border-color: #000000; border-radius: var(--block-radius); background: var(--block-background-fill); width: 100%; line-height: var(--line-sm); } .contour_video { display: flex; flex-direction: column; justify-content: center; align-items: center; z-index: var(--layer-5); border-radius: var(--block-radius); background: var(--background-fill-primary); padding: 0 var(--size-6); max-height: var(--size-screen-h); overflow: hidden; } """ DEFAULT_TEXT = "A person is " 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, "path": path, } 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(data, video_id, width=700, height=700): url = data["url"] start = data["start"] end = data["end"] score = data["score"] text = data["text"] keyid = data["keyid"] babel_id = data["babel_id"] path = data["path"] trim = f"#t={start},{end}" title = f"""Score = {score} Corresponding text: {text} HumanML3D keyid: {keyid} BABEL keyid: {babel_id} AMASS path: {path}""" # class="wrap default svelte-gjihhp hide" #
# width="{width}" height="{height}" video_html = f""" """ return video_html def retrieve_component(retrieve_function, text, splits_choice, nvids, n_component=24): if text == DEFAULT_TEXT or text == "" or text is None: return [None for _ in range(n_component)] # cannot produce more than n_compoenent nvids = min(nvids, n_component) if "Unseen" in splits_choice: splits = ["test"] else: splits = ["train", "val", "test"] datas = retrieve_function(text, splits=splits, nmax=nvids) htmls = [get_video_html(data, idx) for idx, data 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 if not os.path.exists("data"): gdown.download_folder( "https://drive.google.com/drive/folders/1MgPFgHZ28AMd01M1tJ7YW_1-ut3-4j08", use_cookies=False, ) 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") retrieve_and_show = partial(retrieve_component, retrieve_function) with gr.Blocks(css=CSS, theme=theme) as demo: gr.Markdown(WEBSITE) videos = [] with gr.Row(): with gr.Column(scale=3): with gr.Column(scale=2): text = gr.Textbox( placeholder="Type the motion you want to search with a sentence", 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_choice = gr.Radio( ["All motions", "Unseen motions"], label="Gallery of motion", value="All motions", info="The motion gallery is coming from HumanML3D", ) with gr.Column(scale=1): # nvideo_slider = gr.Slider(minimum=4, maximum=24, step=4, value=8, label="Number of videos") nvideo_slider = gr.Radio( [4, 8, 12, 16, 24], label="Videos", value=8, info="Number of videos to display", ) with gr.Column(scale=2): def retrieve_example(text, splits_choice, nvideo_slider): return retrieve_and_show(text, splits_choice, nvideo_slider) examples = gr.Examples( examples=[[x, None, None] for x in EXAMPLES], inputs=[text, splits_choice, nvideo_slider], examples_per_page=20, run_on_click=False, cache_examples=False, fn=retrieve_example, outputs=[], ) i = -1 # should indent for _ in range(6): with gr.Row(): for _ in range(4): i += 1 video = gr.HTML() videos.append(video) # connect the examples to the output # a bit hacky examples.outputs = videos def load_example(example_id): processed_example = examples.non_none_processed_examples[example_id] return gr.utils.resolve_singleton(processed_example) examples.dataset.click( load_example, inputs=[examples.dataset], outputs=examples.inputs_with_examples, # type: ignore show_progress=False, postprocess=False, queue=False, ).then(fn=retrieve_example, inputs=examples.inputs, outputs=videos) btn.click( fn=retrieve_and_show, inputs=[text, splits_choice, nvideo_slider], outputs=videos, ) text.submit( fn=retrieve_and_show, inputs=[text, splits_choice, nvideo_slider], outputs=videos, ) splits_choice.change( fn=retrieve_and_show, inputs=[text, splits_choice, nvideo_slider], outputs=videos, ) nvideo_slider.change( fn=retrieve_and_show, inputs=[text, splits_choice, nvideo_slider], outputs=videos, ) def clear_videos(): return [None for x in range(24)] + [DEFAULT_TEXT] clear.click(fn=clear_videos, outputs=videos + [text]) demo.launch()