import os import gradio as gr import plotly.graph_objects as go import sys import torch from huggingface_hub import hf_hub_download import numpy as np import random os.system("git clone https://github.com/luost26/diffusion-point-cloud") sys.path.append("diffusion-point-cloud") from models.vae_gaussian import * from models.vae_flow import * airplane=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_airplane.pt",revision="main") chair=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_chair.pt",revision="main") device='cuda' if torch.cuda.is_available() else 'cpu' ckpt_airplane = torch.load(airplane,map_location=torch.device(device)) ckpt_chair = torch.load(chair,map_location=torch.device(device)) def seed_all(seed): torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) def normalize_point_clouds(pcs,mode): if mode is None: return pcs for i in range(pcs.size(0)): pc = pcs[i] if mode == 'shape_unit': shift = pc.mean(dim=0).reshape(1, 3) scale = pc.flatten().std().reshape(1, 1) elif mode == 'shape_bbox': pc_max, _ = pc.max(dim=0, keepdim=True) # (1, 3) pc_min, _ = pc.min(dim=0, keepdim=True) # (1, 3) shift = ((pc_min + pc_max) / 2).view(1, 3) scale = (pc_max - pc_min).max().reshape(1, 1) / 2 pc = (pc - shift) / scale pcs[i] = pc return pcs def predict(Seed,ckpt): if Seed==None: Seed=777 seed_all(Seed) if ckpt['args'].model == 'gaussian': model = GaussianVAE(ckpt['args']).to(device) elif ckpt['args'].model == 'flow': model = FlowVAE(ckpt['args']).to(device) model.load_state_dict(ckpt['state_dict']) # Generate Point Clouds gen_pcs = [] with torch.no_grad(): z = torch.randn([1, ckpt['args'].latent_dim]).to(device) x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility) gen_pcs.append(x.detach().cpu()) gen_pcs = torch.cat(gen_pcs, dim=0)[:1] gen_pcs = normalize_point_clouds(gen_pcs, mode="shape_bbox") return gen_pcs[0] def generate(seed,value): if value=="Airplane": ckpt=ckpt_airplane elif value=="Chair": ckpt=ckpt_chair else : ckpt=ckpt_airplane colors=(238, 75, 43) points=predict(seed,ckpt) num_points=points.shape[0] fig = go.Figure( data=[ go.Scatter3d( x=points[:,0], y=points[:,1], z=points[:,2], mode='markers', marker=dict(size=1, color=colors) ) ], layout=dict( scene=dict( xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False) ) ) ) return fig markdown=f''' # Diffusion Probabilistic Models for 3D Point Cloud Generation [[The Paper](https://arxiv.org/abs/2103.01458)] [[Original Code](https://github.com/luost26/diffusion-point-cloud)] The space demo for our CVPR 2021 paper "Diffusion Probabilistic Models for 3D Point Cloud Generation". It is running on {device} ''' with gr.Blocks() as demo: with gr.Column(): with gr.Row(): gr.Markdown(markdown) with gr.Row(): seed = gr.Slider( minimum=0, maximum=2**16,label='Seed') value=gr.Dropdown(choices=["Airplane","Chair"],label="Choose Model Type") btn = gr.Button(value="Generate") point_cloud = gr.Plot() demo.load(generate, [seed,value], point_cloud) btn.click(generate, [seed,value], point_cloud) demo.launch()