Spaces:
Sleeping
Sleeping
File size: 6,934 Bytes
2fd6166 9d94b63 2fd6166 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
"""
Demo built with gradio
"""
import pickle as pkl
import sys, os
import os.path as osp
from typing import Iterable, Optional
from functools import partial
import trimesh
from torch.utils.data import DataLoader
import cv2
from accelerate import Accelerator
from tqdm import tqdm
from glob import glob
sys.path.append(os.getcwd())
import hydra
import torch
import numpy as np
import imageio
import gradio as gr
import plotly.graph_objs as go
import training_utils
from configs.structured import ProjectConfig
from demo import DemoRunner
from dataset.demo_dataset import DemoDataset
md_description="""
# HDM Interaction Reconstruction Demo
### Official Implementation of the paper \"Template Free Reconstruction of Human Object Interaction\", CVPR'24.
[Project Page](https://virtualhumans.mpi-inf.mpg.de/procigen-hdm/)|[Code](https://github.com/xiexh20/HDM)|[Dataset](https://edmond.mpg.de/dataset.xhtml?persistentId=doi:10.17617/3.2VUEUS )|[Paper](https://virtualhumans.mpi-inf.mpg.de/procigen-hdm/paper-lowreso.pdf)
Upload your own human object interaction image and get full 3D reconstruction!
## Citation
```
@inproceedings{xie2023template_free,
title = {Template Free Reconstruction of Human-object Interaction with Procedural Interaction Generation},
author = {Xie, Xianghui and Bhatnagar, Bharat Lal and Lenssen, Jan Eric and Pons-Moll, Gerard},
booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2024},
}
```
"""
def plot_points(colors, coords):
"""
use plotly to visualize 3D point with colors
"""
trace = go.Scatter3d(x=coords[:, 0], y=coords[:, 1], z=coords[:, 2], mode='markers',
marker=dict(
size=2,
color=colors
))
layout = go.Layout(
scene=dict(
xaxis=dict(
title="",
showgrid=False,
zeroline=False,
showline=False,
ticks='',
showticklabels=False
),
yaxis=dict(
title="",
showgrid=False,
zeroline=False,
showline=False,
ticks='',
showticklabels=False
),
zaxis=dict(
title="",
showgrid=False,
zeroline=False,
showline=False,
ticks='',
showticklabels=False
),
),
margin=dict(l=0, r=0, b=0, t=0),
showlegend=False
)
fig = go.Figure(data=[trace], layout=layout)
return fig
def inference(runner: DemoRunner, cfg: ProjectConfig, rgb, mask_hum, mask_obj, std_coverage, input_seed):
"""
given user input, run inference
:param runner:
:param cfg:
:param rgb: (h, w, 3), np array
:param mask_hum: (h, w, 3), np array
:param mask_obj: (h, w, 3), np array
:param std_coverage: float value, used to estimate camera translation
:param input_seed: random seed
:return: path to the 3D reconstruction, and an interactive 3D figure for visualizing the point cloud
"""
# Set random seed
training_utils.set_seed(int(input_seed))
data = DemoDataset([], (cfg.dataset.image_size, cfg.dataset.image_size),
std_coverage)
batch = data.image2batch(rgb, mask_hum, mask_obj)
out_stage1, out_stage2 = runner.forward_batch(batch, cfg)
points = out_stage2.points_packed().cpu().numpy()
colors = out_stage2.features_packed().cpu().numpy()
fig = plot_points(colors, points)
# save tmp point cloud
outdir = './results'
os.makedirs(outdir, exist_ok=True)
trimesh.PointCloud(points, colors).export(outdir + f"/pred_std{std_coverage}_seed{input_seed}_stage2.ply")
trimesh.PointCloud(out_stage1.points_packed().cpu().numpy(),
out_stage1.features_packed().cpu().numpy()).export(outdir + f"/pred_std{std_coverage}_seed{input_seed}_stage1.ply")
return fig, outdir + f"/pred_std{std_coverage}_seed{input_seed}_stage2.ply"
@hydra.main(config_path='configs', config_name='configs', version_base='1.1')
def main(cfg: ProjectConfig):
# Setup model
runner = DemoRunner(cfg)
# runner = None # without model initialization, it shows one line of thumbnail
# Setup interface
demo = gr.Blocks(title="HDM Interaction Reconstruction Demo")
with demo:
gr.Markdown(md_description)
gr.HTML("""<h1 style="text-align:center; color:#10768c">HDM Demo</h1>""")
gr.HTML("""<h3 style="text-align:center; color:#10768c">Instruction: Upload RGB, human, object masks and then click reconstruct.</h1>""")
# Input data
with gr.Row():
input_rgb = gr.Image(label='Input RGB', type='numpy')
input_mask_hum = gr.Image(label='Human mask', type='numpy')
with gr.Row():
input_mask_obj = gr.Image(label='Object mask', type='numpy')
with gr.Column():
# TODO: add hint for this value here
input_std = gr.Number(label='Gaussian std coverage', value=3.5)
input_seed = gr.Number(label='Random seed', value=42)
# Output visualization
with gr.Row():
pc_plot = gr.Plot(label="Reconstructed point cloud")
out_pc_download = gr.File(label="3D reconstruction for download") # this allows downloading
gr.HTML("""<br/>""")
# Control
with gr.Row():
button_recon = gr.Button("Start Reconstruction", interactive=True, variant='secondary')
button_recon.click(fn=partial(inference, runner, cfg),
inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed],
outputs=[pc_plot, out_pc_download])
gr.HTML("""<br/>""")
# Example input
example_dir = cfg.run.code_dir_abs+"/examples"
rgb, ps, obj = 'k1.color.jpg', 'k1.person_mask.png', 'k1.obj_rend_mask.png'
example_images = gr.Examples([
[f"{example_dir}/017450/{rgb}", f"{example_dir}/017450/{ps}", f"{example_dir}/017450/{obj}", 3.0, 42],
[f"{example_dir}/002446/{rgb}", f"{example_dir}/002446/{ps}", f"{example_dir}/002446/{obj}", 3.0, 42],
[f"{example_dir}/053431/{rgb}", f"{example_dir}/053431/{ps}", f"{example_dir}/053431/{obj}", 3.8, 42],
[f"{example_dir}/158107/{rgb}", f"{example_dir}/158107/{ps}", f"{example_dir}/158107/{obj}", 3.8, 42],
], inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed],)
# demo.launch(share=True)
# Enabling queue for runtime>60s, see: https://github.com/tloen/alpaca-lora/issues/60#issuecomment-1510006062
demo.queue(concurrency_count=3).launch(share=True)
if __name__ == '__main__':
main() |