scatter in plotly

app.py

@@ -2,7 +2,7 @@ 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, create_3d_pc
+from utils import create_3d_obj, create_3d_pc, point_cloud
 import numpy as np
 
 
@@ -13,9 +13,9 @@ def snap(image, video):
     
     segment_predictor = SegmentPredictor()
     sam_result = segment_predictor.predict(image)
-    seg_gltf_path = create_3d_obj(np.array(sam_result), depth_result, path='./seg.gltf')
+    fig = point_cloud(np.array(sam_result), depth_result)
 
-    return [image, depth_result, sam_result, rgb_gltf_path, seg_gltf_path]#[depth_result, gltf_path, gltf_path]
+    return [image, depth_result, sam_result, rgb_gltf_path, fig]#[depth_result, gltf_path, gltf_path]
 
 
 demo = gr.Interface(
@@ -27,8 +27,7 @@ demo = gr.Interface(
              gr.Image(label="predicted segmentation"),
              gr.Model3D(label="3D mesh reconstruction - RGB",
                         clear_color=[1.0, 1.0, 1.0, 1.0]),
-             gr.Model3D(label="3D mesh reconstruction - Segmented",
-                        clear_color=[1.0, 1.0, 1.0, 1.0])]
+             gr.Plot()]
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -6,4 +6,7 @@ torch
 torchvision
 opencv-python
 transformers
-open3d
+open3d
+plotly
+pandas
+numpy

utils.py

@@ -1,5 +1,9 @@
 import numpy as np
 import open3d as o3d
+import open3d as o3d
+import plotly.express as px
+import numpy as np
+import pandas as pd
 
 
 def create_3d_obj(rgb_image, depth_image, depth=10, path='./image.gltf'):
@@ -85,3 +89,21 @@ def create_3d_pc(rgb_image, depth_image, depth=10):
     o3d.io.write_point_cloud(filename, pcd)
 
     return filename # Return the file path where the PLY file is saved
+
+
+def point_cloud(rgb_image, depth_image):
+    # Step 2: Create an RGBD image from the RGB and depth images
+    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb_image, depth_image, convert_rgb_to_intensity=False)
+    # Step 3: Create a PointCloud from the RGBD image
+    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault))
+    # Step 4: Convert PointCloud data to a NumPy array
+    points = np.asarray(pcd.points)
+    colors = np.asarray(pcd.colors)
+    # Step 5: Create a DataFrame from the NumPy arrays
+    data = {'x': points[:, 0], 'y': points[:, 1], 'z': points[:, 2],
+            'red': colors[:, 0], 'green': colors[:, 1], 'blue': colors[:, 2]}
+    df = pd.DataFrame(data)
+    # Step 6: Create a 3D scatter plot using Plotly Express
+    fig = px.scatter_3d(df, x='x', y='y', z='z', color='red', size_max=0.1)
+    
+    return fig