SerdarHelli/diffusion-point-cloud

1import os2import gradio as gr3import plotly.graph_objects as go4import sys5import torch6from huggingface_hub import hf_hub_download7import numpy as np8import random9 10os.system("git clone https://github.com/luost26/diffusion-point-cloud")11sys.path.append("diffusion-point-cloud")12 13#Codes reference : https://github.com/luost26/diffusion-point-cloud14 15from models.vae_gaussian import *16from models.vae_flow import *17 18airplane=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_airplane.pt",revision="main")19chair="./GEN_chair.pt"20 21device='cuda' if torch.cuda.is_available() else 'cpu'22 23ckpt_airplane = torch.load(airplane,map_location=torch.device(device))24ckpt_chair = torch.load(chair,map_location=torch.device(device))25 26def seed_all(seed):27    torch.manual_seed(seed)28    np.random.seed(seed)29    random.seed(seed)30 31def normalize_point_clouds(pcs,mode):32    if mode is None:33        return pcs34    for i in range(pcs.size(0)):35        pc = pcs[i]36        if mode == 'shape_unit':37            shift = pc.mean(dim=0).reshape(1, 3)38            scale = pc.flatten().std().reshape(1, 1)39        elif mode == 'shape_bbox':40            pc_max, _ = pc.max(dim=0, keepdim=True) # (1, 3)41            pc_min, _ = pc.min(dim=0, keepdim=True) # (1, 3)42            shift = ((pc_min + pc_max) / 2).view(1, 3)43            scale = (pc_max - pc_min).max().reshape(1, 1) / 244        pc = (pc - shift) / scale45        pcs[i] = pc46    return pcs47 48    49 50 51def predict(Seed,ckpt):52  if Seed==None:53    Seed=777 54  seed_all(Seed)55 56  if ckpt['args'].model == 'gaussian':57      model = GaussianVAE(ckpt['args']).to(device)58  elif ckpt['args'].model == 'flow':59      model = FlowVAE(ckpt['args']).to(device)60 61  model.load_state_dict(ckpt['state_dict'])62  # Generate Point Clouds63  gen_pcs = []64  with torch.no_grad():65      z = torch.randn([1, ckpt['args'].latent_dim]).to(device)66      x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility)67      gen_pcs.append(x.detach().cpu())68  gen_pcs = torch.cat(gen_pcs, dim=0)[:1]69  gen_pcs = normalize_point_clouds(gen_pcs, mode="shape_bbox")70 71  return gen_pcs[0]72 73def generate(seed,value):74    if value=="Airplane":75      ckpt=ckpt_airplane76    elif value=="Chair":77      ckpt=ckpt_chair78    else :79       ckpt=ckpt_airplane80 81    colors=(238, 75, 43)82    points=predict(seed,ckpt)83    num_points=points.shape[0]84 85 86    fig = go.Figure(87        data=[88            go.Scatter3d(89                x=points[:,0], y=points[:,1], z=points[:,2], 90                mode='markers',91                marker=dict(size=1, color=colors)92            )93        ],94        layout=dict(95            scene=dict(96                xaxis=dict(visible=False),97                yaxis=dict(visible=False),98                zaxis=dict(visible=False)99            )100        )101    )102    return fig103    104markdown=f'''105  # Diffusion Probabilistic Models for 3D Point Cloud Generation106 107  108  [The space demo for the CVPR 2021 paper "Diffusion Probabilistic Models for 3D Point Cloud Generation".](https://arxiv.org/abs/2103.01458)109  110  [For the official implementation.](https://github.com/luost26/diffusion-point-cloud)111 112  ### Future Work based on interest113  - Adding new models for new type objects114  - New Customization 115  116  117 118  It is running on {device}119  120 121'''122with gr.Blocks() as demo:123    with gr.Column():124        with gr.Row():125            gr.Markdown(markdown)126        with gr.Row():127            seed = gr.Slider( minimum=0, maximum=2**16,label='Seed')128            value=gr.Dropdown(choices=["Airplane","Chair"],label="Choose Model Type")129            #truncate_std = gr.Slider( minimum=1, maximum=2,label='Truncate Std')130 131        btn = gr.Button(value="Generate")132        point_cloud = gr.Plot()133    demo.load(generate, [seed,value], point_cloud)134    btn.click(generate, [seed,value], point_cloud)135 136demo.launch()