Update app.py

app.py

@@ -43,42 +43,47 @@ def process_image(image_path):
     else:
         depth_image = np.zeros_like(output, dtype='uint8')  # Handle empty output
     
-    try:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
-    except Exception:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path, depth=8)
+    gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
     
-    return Image.fromarray(depth_image), gltf_path, gltf_path
+    if gltf_path and Path(gltf_path).exists():
+        return Image.fromarray(depth_image), gltf_path, gltf_path
+    else:
+        return Image.fromarray(depth_image), None, "3D model generation failed"
 
-def create_3d_obj(rgb_image, depth_image, image_path, depth=10):
-    depth_o3d = o3d.geometry.Image(depth_image)
-    image_o3d = o3d.geometry.Image(rgb_image)
-    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
-        image_o3d, depth_o3d, convert_rgb_to_intensity=False)
+def create_3d_obj(rgb_image, depth_image, image_path):
+    try:
+        depth_o3d = o3d.geometry.Image(depth_image)
+        image_o3d = o3d.geometry.Image(rgb_image)
+        rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
+            image_o3d, depth_o3d, convert_rgb_to_intensity=False)
 
-    w, h = depth_image.shape[1], depth_image.shape[0]
-    camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
-    camera_intrinsic.set_intrinsics(w, h, 500, 500, w / 2, h / 2)
+        w, h = depth_image.shape[1], depth_image.shape[0]
+        camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
+        camera_intrinsic.set_intrinsics(w, h, 500, 500, w / 2, h / 2)
 
-    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)
-    pcd.estimate_normals(
-        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
-    pcd.orient_normals_towards_camera_location(camera_location=np.array([0., 0., 1000.]))
+        pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)
+        pcd.estimate_normals(
+            search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
+        pcd.orient_normals_towards_camera_location(camera_location=np.array([0., 0., 1000.]))
 
-    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug):
         mesh_raw, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-            pcd, depth=depth, width=0, scale=1.1, linear_fit=True)
+            pcd, depth=10, width=0, scale=1.1, linear_fit=True)
 
-    voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / 256
-    mesh = mesh_raw.simplify_vertex_clustering(voxel_size=voxel_size)
+        if not mesh_raw.has_triangles():
+            return None  # Mesh generation failed
+        
+        # Center the mesh for better preview
+        bbox = pcd.get_axis_aligned_bounding_box()
+        mesh_raw.translate(-bbox.get_center())
 
-    bbox = pcd.get_axis_aligned_bounding_box()
-    mesh_crop = mesh.crop(bbox)
-    
-    gltf_path = f'./{image_path.stem}.gltf'
-    o3d.io.write_triangle_mesh(gltf_path, mesh_crop, write_triangle_uvs=True)
-    
-    return gltf_path
+        # Save the 3D model
+        gltf_path = f'./{image_path.stem}.gltf'
+        o3d.io.write_triangle_mesh(gltf_path, mesh_raw, write_triangle_uvs=True)
+        return gltf_path
+
+    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."
@@ -109,3 +114,4 @@ if __name__ == "__main__":
 
 
 
+