SerdarHelli commited on
Commit
0d0e451
1 Parent(s): c15a10c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +114 -0
app.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import plotly.graph_objects as go
4
+ import sys
5
+ import torch
6
+ from huggingface_hub import hf_hub_download
7
+
8
+ os.system("git clone https://github.com/luost26/diffusion-point-cloud")
9
+ sys.path.append("diffusion-point-cloud")
10
+
11
+
12
+ from models.vae_gaussian import *
13
+ from models.vae_flow import *
14
+
15
+ airplane=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_airplane.pt",revision="main")
16
+ chair=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_chair.pt",revision="main")
17
+
18
+
19
+ ckpt_airplane = torch.load(airplane)
20
+ ckpt_chair = torch.load(chair)
21
+
22
+ def normalize_point_clouds(pcs,mode):
23
+ if mode is None:
24
+ return pcs
25
+ for i in range(pcs.size(0)):
26
+ pc = pcs[i]
27
+ if mode == 'shape_unit':
28
+ shift = pc.mean(dim=0).reshape(1, 3)
29
+ scale = pc.flatten().std().reshape(1, 1)
30
+ elif mode == 'shape_bbox':
31
+ pc_max, _ = pc.max(dim=0, keepdim=True) # (1, 3)
32
+ pc_min, _ = pc.min(dim=0, keepdim=True) # (1, 3)
33
+ shift = ((pc_min + pc_max) / 2).view(1, 3)
34
+ scale = (pc_max - pc_min).max().reshape(1, 1) / 2
35
+ pc = (pc - shift) / scale
36
+ pcs[i] = pc
37
+ return pcs
38
+
39
+ def predict(Seed,ckpt):
40
+ if Seed==None:
41
+ Seed=777
42
+ seed_all(Seed)
43
+
44
+ if ckpt['args'].model == 'gaussian':
45
+ model = GaussianVAE(ckpt['args']).to("cuda")
46
+ elif ckpt['args'].model == 'flow':
47
+ model = FlowVAE(ckpt['args']).to("cuda")
48
+
49
+ model.load_state_dict(ckpt['state_dict'])
50
+ # Generate Point Clouds
51
+ gen_pcs = []
52
+ with torch.no_grad():
53
+ z = torch.randn([1, ckpt['args'].latent_dim]).to("cuda")
54
+ x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility)
55
+ gen_pcs.append(x.detach().cpu())
56
+ gen_pcs = torch.cat(gen_pcs, dim=0)[:1]
57
+ gen_pcs = normalize_point_clouds(gen_pcs, mode="shape_bbox")
58
+
59
+ return gen_pcs[0]
60
+
61
+ def generate(seed,value):
62
+ if value=="Airplane":
63
+ ckpt=ckpt_airplane
64
+ elif value=="Chair":
65
+ ckpt=ckpt_chair
66
+ else :
67
+ ckpt=ckpt_airplane
68
+
69
+ print(value)
70
+ colors=(238, 75, 43)
71
+ points=predict(seed,ckpt)
72
+ num_points=points.shape[0]
73
+
74
+
75
+ fig = go.Figure(
76
+ data=[
77
+ go.Scatter3d(
78
+ x=points[:,0], y=points[:,1], z=points[:,2],
79
+ mode='markers',
80
+ marker=dict(size=1, color=colors)
81
+ )
82
+ ],
83
+ layout=dict(
84
+ scene=dict(
85
+ xaxis=dict(visible=False),
86
+ yaxis=dict(visible=False),
87
+ zaxis=dict(visible=False)
88
+ )
89
+ )
90
+ )
91
+ return fig
92
+ markdown=f'''
93
+ # Diffusion Probabilistic Models for 3D Point Cloud Generation
94
+
95
+ [[The Paper](https://arxiv.org/abs/2103.01458)] [[Original Code](https://github.com/luost26/diffusion-point-cloud)]
96
+
97
+ The space demo for our CVPR 2021 paper "Diffusion Probabilistic Models for 3D Point Cloud Generation".
98
+
99
+
100
+ '''
101
+ with gr.Blocks() as demo:
102
+ with gr.Column():
103
+ with gr.Row():
104
+ gr.Markdown(markdown)
105
+ with gr.Row():
106
+ seed = gr.Slider( minimum=0, maximum=2**16,label='Seed')
107
+ value=gr.Dropdown(choices=["Airplane","Chair"],label="Choose Model Type")
108
+
109
+ btn = gr.Button(value="Generate")
110
+ point_cloud = gr.Plot()
111
+ demo.load(generate, [seed,value], point_cloud)
112
+ btn.click(generate, [seed,value], point_cloud)
113
+
114
+ demo.launch()