Update app.py

app.py

@@ -24,6 +24,7 @@ import numpy as np
 import open3d as o3d
 import open3d.visualization.rendering as rendering
 
+import plotly.graph_objects as go
 from PIL import Image
 from tqdm.auto import tqdm
 from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
@@ -174,15 +175,39 @@ def compose_pointe(prompt, weights, version):
     prompt_list = [x.strip() for x in prompt.split("|")]
     pcd = generate_pcd(prompt_list)
     pc = generate_fig(pcd)
-    mesh = generate_mesh(pc)
-    timestr = time.strftime("%Y%m%d-%H%M%S")
-    mesh_path = os.path.join(f'{timestr}.ply')
-    with open(mesh_path, 'wb') as f:
-        mesh.write_ply(f)
-    image_frames = generate_video(mesh_path)
-    gif_path = os.path.join(f'{timestr}.gif')
-    image_frames[0].save(gif_path, save_all=True, optimizer=False, duration=5, append_images=image_frames[1:], loop=0)
-    return f'{timestr}.gif'
+
+    fig = go.Figure(
+        data=[
+            go.Scatter3d(
+                x=pc.coords[:, 0], y=pc.coords[:, 1], z=pc.coords[:, 2],
+                mode='markers',
+                marker=dict(
+                    size=2,
+                    color=['rgb({},{},{})'.format(r, g, b) for r, g, b in
+                           zip(pc.channels["R"], pc.channels["G"], pc.channels["B"])],
+                )
+            )
+        ],
+        layout=dict(
+            scene=dict(
+                xaxis=dict(visible=False),
+                yaxis=dict(visible=False),
+                zaxis=dict(visible=False)
+            )
+        ),
+    )
+    return fig
+
+    # huggingface failed to render, so we only visualize pointclouds
+    # mesh = generate_mesh(pc)
+    # timestr = time.strftime("%Y%m%d-%H%M%S")
+    # mesh_path = os.path.join(f'{timestr}.ply')
+    # with open(mesh_path, 'wb') as f:
+    #     mesh.write_ply(f)
+    # image_frames = generate_video(mesh_path)
+    # gif_path = os.path.join(f'{timestr}.gif')
+    # image_frames[0].save(gif_path, save_all=True, optimizer=False, duration=5, append_images=image_frames[1:], loop=0)
+    # return f'{timestr}.gif'
 
 
 def compose_clevr_objects(prompt, weights, steps):
@@ -353,7 +378,8 @@ with gr.Blocks() as demo:
                 asset_weights = gr.Textbox(value="7.5 | 7.5", label="Weights")
             with gr.Row():
                 asset_model = gr.Radio(['Point-E'], type="value", label='Text to 3D model', value='Point-E')
-            asset_output = gr.Image(label='GIF')
+            # asset_output = gr.Image(label='GIF')
+            asset_output = gr.Plot(label='Plot')
             asset_button = gr.Button("Generate")
             asset_examples = gr.Examples(examples=pointe_examples, inputs=[asset_input, asset_weights, asset_model])