Instruction:
More example results can be found in our Project Page.
Have fun!
""" 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, input_cls, input_scheduler): """ 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 :param input_cls: the object category of the input image :param input_scheduler: reverse sampling scheduler, ddim or ddpm :return: path to the 3D reconstruction, and an interactive 3D figure for visualizing the point cloud """ log = "" try: # 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) if input_cls != 'general': log += f"Reloading fine-tuned checkpoint of category {input_cls}\n" runner.reload_checkpoint(input_cls) cfg.run.diffusion_scheduler = input_scheduler cfg.run.num_inference_steps = 1000 if input_scheduler == 'ddpm' else 100 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_{input_cls}.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_{input_cls}.ply") log += 'Successfully reconstructed the image.' outfile = outdir + f"/pred_std{std_coverage}_seed{input_seed}_stage2_{input_cls}.ply" except Exception as e: log = traceback.format_exc() fig, outfile = None, None return fig, outfile, log @hydra.main(config_path='configs', config_name='configs', version_base='1.1') def main(cfg: ProjectConfig): # Setup model runner = DemoRunner(cfg) # Setup interface demo = gr.Blocks(title="HDM Interaction Reconstruction Demo") with demo: gr.HTML(html_str) # 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(): input_std = gr.Number(label='Gaussian std coverage', value=3.5, info="This value is used to estimate camera translation to project the points." "The larger value, the camera is farther away. It is category-dependent. " "We empirically found these values are suitable: backpack-3.5, ball-3.0, bottle-3.0," "box-3.5, chair-3.8, skateboard-3.0, suitcase-3.2, table-3.5. " "If you are not sure, 3.5 is a good start point.") input_cls = gr.Dropdown(label='Object category', info='We fine tuned the model for some specific categories. ' 'Reconstructing using these models should lead to better result ' 'for these specific categories. Simply select the category that ' 'fits the object from input image.', choices=['general', 'backpack', 'ball', 'bottle', 'box', 'chair', 'skateboard', 'suitcase', 'table'], value='general') input_seed = gr.Number(label='Random seed', value=42) input_scheduler = gr.Dropdown(label='Diffusion scheduler', info='Reverse diffusion scheduler: DDIM is 10x faster', choices=['ddpm', 'ddim'], value='ddim') # Output visualization with gr.Row(): pc_plot = gr.Plot(label="Reconstructed point cloud") out_pc_download = gr.File(label="Download results") # this allows downloading with gr.Row(): out_log = gr.TextArea(label='Output log') gr.HTML("""