feat: add PLY to GLB conversion and update UI for dual display and download

app.py

@@ -16,6 +16,8 @@ subprocess.run(
 
 import gradio as gr  # import AFTER the pip install above
 from huggingface_hub import HfApi
+import trimesh
+from plyfile import PlyData
 
 # install custom wheels for gaussian splatting
 subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"))
@@ -57,6 +59,43 @@ def get_dust3r_args_parser():
     return parser
 
 
+def convert_ply_to_glb(ply_path, glb_path):
+    try:
+        plydata = PlyData.read(ply_path)
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])), axis=1)
+        
+        # Extract DC features (colors)
+        # f_dc_0, f_dc_1, f_dc_2 are SH coefficients for R, G, B
+        f_dc_0 = np.asarray(plydata.elements[0]["f_dc_0"])
+        f_dc_1 = np.asarray(plydata.elements[0]["f_dc_1"])
+        f_dc_2 = np.asarray(plydata.elements[0]["f_dc_2"])
+        
+        # SH2RGB: sh * C0 + 0.5
+        # C0 = 0.28209479177387814
+        C0 = 0.28209479177387814
+        r = f_dc_0 * C0 + 0.5
+        g = f_dc_1 * C0 + 0.5
+        b = f_dc_2 * C0 + 0.5
+        
+        colors = np.stack((r, g, b), axis=1)
+        # Clip to [0, 1]
+        colors = np.clip(colors, 0, 1)
+        # Convert to uint8 for trimesh
+        colors = (colors * 255).astype(np.uint8)
+        
+        # Create PointCloud
+        pcd = trimesh.points.PointCloud(vertices=xyz, colors=colors)
+        
+        # Export
+        pcd.export(glb_path)
+        return True
+    except Exception as e:
+        print(f"Error converting PLY to GLB: {e}")
+        return False
+
+
 @spaces.GPU(duration=150)
 def process(inputfiles, input_path=None):
 
@@ -219,11 +258,18 @@ def process(inputfiles, input_path=None):
         print("Failed to upload PLY to hub:", e)
         ply_url = f"LOCAL:{output_ply_path}"
 
+    # Convert PLY to GLB for visualization
+    output_glb_path = output_ply_path.replace('.ply', '.glb')
+    if not convert_ply_to_glb(output_ply_path, output_glb_path):
+        output_glb_path = None
+
     # return:
     #  1) video path (for gr.Video)
     #  2) ply URL (for API + textbox)
-    #  3) local ply path (for gr.Model3D viewer)
-    return output_video_path, ply_url, output_ply_path
+    #  3) ply file path (for gr.File download)
+    #  4) ply file path (for gr.Model3D viewer)
+    #  5) glb file path (for gr.Model3D viewer)
+    return output_video_path, ply_url, output_ply_path, output_ply_path, output_glb_path
     ##################################################################################################################################################
 
 
@@ -269,8 +315,13 @@ with block:
         with gr.Tab("Output"):
             with gr.Column(scale=2):
                 with gr.Group():
-                    output_model = gr.Model3D(
-                        label="3D Dense Model under Gaussian Splats Formats, need more time to visualize",
+                    output_model_glb = gr.Model3D(
+                        label="3D Model (GLB Point Cloud)",
+                        interactive=False,
+                        camera_position=[0.5, 0.5, 1],
+                    )
+                    output_model_ply = gr.Model3D(
+                        label="Original PLY (Gaussian Splat)",
                         interactive=False,
                         camera_position=[0.5, 0.5, 1],
                     )
@@ -285,17 +336,21 @@ with block:
                     label="PLY download URL",
                     interactive=False,
                 )
+                output_download = gr.File(
+                    label="Download PLY",
+                    interactive=False,
+                )
             with gr.Column(scale=1):
                 output_video = gr.Video(label="video")
                 
-    button_gen.click(process, inputs=[inputfiles], outputs=[output_video, output_file, output_model])
+    button_gen.click(process, inputs=[inputfiles], outputs=[output_video, output_file, output_download, output_model_ply, output_model_glb])
     
     gr.Examples(
         examples=[
             "sora-santorini-3-views",
         ],
         inputs=[input_path],
-        outputs=[output_video, output_file, output_model],
+        outputs=[output_video, output_file, output_download, output_model_ply, output_model_glb],
         fn=lambda x: process(inputfiles=None, input_path=x),
         cache_examples=True,
         label='Sparse-view Examples'