Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from PIL import Image
 import open3d as o3d
 from pathlib import Path
+import subprocess
 
 # Load model and feature extractor
 feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
@@ -43,10 +44,10 @@ def process_image(image_path):
     else:
         depth_image = np.zeros_like(output, dtype='uint8')  # Handle empty output
     
-    gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
+    glb_path = create_3d_obj(np.array(image), depth_image, image_path)
     
-    if gltf_path and Path(gltf_path).exists():
-        return Image.fromarray(depth_image), gltf_path, gltf_path
+    if glb_path and Path(glb_path).exists():
+        return Image.fromarray(depth_image), glb_path, glb_path
     else:
         return Image.fromarray(depth_image), None, "3D model generation failed"
 
@@ -70,23 +71,33 @@ def create_3d_obj(rgb_image, depth_image, image_path):
             pcd, depth=10, width=0, scale=1.1, linear_fit=True)
 
         if not mesh_raw.has_triangles():
+            print("Mesh generation failed: No triangles in mesh")
             return None  # Mesh generation failed
         
         # Center the mesh for better preview
         bbox = pcd.get_axis_aligned_bounding_box()
         mesh_raw.translate(-bbox.get_center())
 
-        # Save the 3D model
-        gltf_path = f'./{image_path.stem}.gltf'
+        # Save the 3D model as .gltf
+        gltf_path = str(Path.cwd() / f"{image_path.stem}.gltf")
         o3d.io.write_triangle_mesh(gltf_path, mesh_raw, write_triangle_uvs=True)
-        return gltf_path
+
+        # Convert .gltf to .glb
+        glb_path = gltf_path.replace(".gltf", ".glb")
+        subprocess.run(["npx", "gltf-pipeline", "-i", gltf_path, "-o", glb_path])
+
+        if Path(glb_path).exists():
+            return glb_path
+        else:
+            print("GLB conversion failed.")
+            return None
 
     except Exception as e:
         print(f"3D model generation failed: {e}")
         return None
 
 title = "Zero-shot Depth Estimation with DPT + 3D Model Preview"
-description = "Upload an image to generate a depth map and reconstruct a 3D model in .gltf format."
+description = "Upload an image to generate a depth map and reconstruct a 3D model in .glb format."
 
 with gr.Blocks() as demo:
     gr.Markdown(f"## {title}")
@@ -99,8 +110,8 @@ with gr.Blocks() as demo:
         
         with gr.Column(scale=2):
             depth_output = gr.Image(label="Predicted Depth", type="pil")
-            model_output = gr.Model3D(label="3D Model Preview (GLTF)")
-            file_output = gr.File(label="Download 3D GLTF File")
+            model_output = gr.Model3D(label="3D Model Preview (GLB)")
+            file_output = gr.File(label="Download 3D GLB File")
 
     generate_button.click(fn=process_image, inputs=[image_input], outputs=[depth_output, model_output, file_output])