examples and voxels version

README.md

@@ -1,6 +1,6 @@
 ---
 title: Dpt Depth Estimation + 3D Voxels
-emoji: ⚡
+emoji: 🧊
 colorFrom: blue
 colorTo: red
 sdk: gradio

app.py

@@ -1,3 +1,4 @@
+from email.policy import default
 import gradio as gr
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 import torch
@@ -11,7 +12,8 @@ feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
 model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
 
 
-def process_image(image_path):
+def process_image(image_path, voxel_s):
+    voxel_s = max(voxel_s/500, 0.0001)
     image_path = Path(image_path)
     image_raw = Image.open(image_path)
     image = image_raw.resize(
@@ -36,20 +38,16 @@ def process_image(image_path):
     output = prediction.cpu().numpy()
     depth_image = (output * 255 / np.max(output)).astype('uint8')
     try:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
+        gltf_path = create_3d_voxels_obj(
+            np.array(image), depth_image, image_path, voxel_s)
         img = Image.fromarray(depth_image)
         return [img, gltf_path, gltf_path]
     except Exception as e:
-        gltf_path = create_3d_obj(
-            np.array(image), depth_image, image_path, depth=8)
-        img = Image.fromarray(depth_image)
-        return [img, gltf_path, gltf_path]
-    except:
         print("Error reconstructing 3D model")
         raise Exception("Error reconstructing 3D model")
 
 
-def create_3d_obj(rgb_image, depth_image, image_path, depth=10):
+def create_3d_voxels_obj(rgb_image, depth_image, image_path, voxel_s):
     depth_o3d = o3d.geometry.Image(depth_image)
     image_o3d = o3d.geometry.Image(rgb_image)
     rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
@@ -79,38 +77,46 @@ def create_3d_obj(rgb_image, depth_image, image_path, depth=10):
                    [0, 0, 1, 0],
                    [0, 0, 0, 1]])
 
-    print('run Poisson surface reconstruction')
-    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
-        mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-            pcd, depth=depth, width=0, scale=1.1, linear_fit=True)
-
-    voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / 256
-    print(f'voxel_size = {voxel_size:e}')
-    mesh = mesh_raw.simplify_vertex_clustering(
-        voxel_size=voxel_size,
-        contraction=o3d.geometry.SimplificationContraction.Average)
-
-    # vertices_to_remove = densities < np.quantile(densities, 0.001)
-    # mesh.remove_vertices_by_mask(vertices_to_remove)
-    bbox = pcd.get_axis_aligned_bounding_box()
-    mesh_crop = mesh.crop(bbox)
+    print('voxels')
+
+    # ref https://towardsdatascience.com/how-to-automate-voxel-modelling-of-3d-point-cloud-with-python-459f4d43a227
+    voxel_size = round(
+        max(pcd.get_max_bound()-pcd.get_min_bound())*voxel_s, 10)
+    print("Voxel size", voxel_size, "voxel_s", voxel_s)
+    voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(
+        pcd, voxel_size=voxel_size)
+    voxels = voxel_grid.get_voxels()
+
+    vox_mesh = o3d.geometry.TriangleMesh()
+    for v in voxels:
+        cube = o3d.geometry.TriangleMesh.create_box(width=1, height=1, depth=1)
+        cube.paint_uniform_color(v.color)
+        cube.translate(v.grid_index, relative=False)
+        vox_mesh += cube
+    print(voxel_grid, vox_mesh)
+
     gltf_path = f'./{image_path.stem}.gltf'
-    o3d.io.write_triangle_mesh(
-        gltf_path, mesh_crop, write_triangle_uvs=True)
+    o3d.io.write_triangle_mesh(gltf_path, vox_mesh, write_triangle_uvs=True)
     return gltf_path
 
 
-title = "Demo: zero-shot depth estimation with DPT + 3D Point Cloud"
-description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object."
-examples = [["examples/" + img] for img in os.listdir("examples/")]
+title = "Demo: zero-shot depth estimation with DPT + 3D Voxels reconstruction"
+description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then reconstruct the 3D model as voxels."
+examples = [["examples/" + img, 10] for img in os.listdir("examples/")]
 
 iface = gr.Interface(fn=process_image,
-                     inputs=[gr.inputs.Image(
-                         type="filepath", label="Input Image")],
-                     outputs=[gr.outputs.Image(label="predicted depth", type="pil"),
-                              gr.outputs.Image3D(label="3d mesh reconstruction", clear_color=[
-                                                 1.0, 1.0, 1.0, 1.0]),
-                              gr.outputs.File(label="3d gLTF")],
+                     inputs=[
+                         gr.inputs.Image(
+                             type="filepath", label="Input Image"),
+                         gr.inputs.Slider(
+                             5, 100, step=1, label="Voxel Size", default=10)
+                     ],
+                     outputs=[
+                         gr.outputs.Image(label="predicted depth", type="pil"),
+                         gr.outputs.Image3D(label="3d mesh reconstruction", clear_color=[
+                             1.0, 1.0, 1.0, 1.0]),
+                         gr.outputs.File(label="3d gLTF")
+                     ],
                      title=title,
                      description=description,
                      examples=examples,