3d

app.py

@@ -2,15 +2,17 @@ import gradio as gr
 from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 import supervision as sv
 from inference import DepthPredictor, SegmentPredictor
-
+from utils import create_3d_obj
+import numpy as np
 
 
 def snap(image, video):
     depth_predictor = DepthPredictor()
     depth_result = depth_predictor.predict(image)
+    gltf_path = create_3d_obj(np.array(image), depth_result)
     #segment_predictor = SegmentPredictor()
     #sam_result = segment_predictor.predict(image)
-    return [depth_result, video]
+    return [depth_result, gltf_path, gltf_path]
 
 
 demo = gr.Interface(

requirements.txt

@@ -5,4 +5,5 @@ supervision
 torch
 torchvision
 opencv-python
-transformers
+transformers
+open3d

utils.py

@@ -0,0 +1,53 @@
+import numpy as np
+import open3d as o3d
+
+
+def create_3d_obj(rgb_image, depth_image, 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)
+    w = int(depth_image.shape[1])
+    h = int(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)
+
+    print('normals')
+    pcd.normals = o3d.utility.Vector3dVector(
+        np.zeros((1, 3)))  # invalidate existing normals
+    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.transform([[1, 0, 0, 0],
+                   [0, -1, 0, 0],
+                   [0, 0, -1, 0],
+                   [0, 0, 0, 1]])
+    pcd.transform([[-1, 0, 0, 0],
+                   [0, 1, 0, 0],
+                   [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)
+    gltf_path = './image.gltf'
+    o3d.io.write_triangle_mesh(
+        gltf_path, mesh_crop, write_triangle_uvs=True)
+    return gltf_path