fix

app.py

@@ -4,12 +4,13 @@ import numpy as np
 import cv2 
 from PIL import Image
 import torch
-from inference import SegmentPredictor
+from inference import SegmentPredictor, DepthPredictor
 from utils import generate_PCL, PCL3, point_cloud
 
 
 
 sam = SegmentPredictor()
+dpt = DepthPredictor()
 red = (255,0,0)
 blue = (0,0,255)
 annos = []
@@ -52,7 +53,7 @@ with block:
         print("depth reconstruction")
         image = inputs[raw_image]
         # depth reconstruction
-        fig = point_cloud(image)
+        fig = dpt.generate_fig(image)
         return {pcl_figure: fig}
     
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)

inference.py

@@ -6,6 +6,9 @@ import torch
 import numpy as np
 from PIL import Image
 import requests
+import open3d as o3d
+import pandas as pd
+import plotly.express as px
 
 class DepthPredictor:
     def __init__(self):
@@ -17,7 +20,7 @@ class DepthPredictor:
     def predict(self, image):
         # prepare image for the model
         encoding = self.feature_extractor(image, return_tensors="pt")
-        
+        self.img = image
         # forward pass
         with torch.no_grad():
             outputs = self.model(**encoding)
@@ -36,6 +39,31 @@ class DepthPredictor:
         #img = Image.fromarray(formatted)
         return formatted
     
+    def generate_pcl(self, image):
+        depth = self.predict(image)
+        # Step 2: Create an RGBD image from the RGB and depth image
+        depth_o3d = o3d.geometry.Image(depth)
+        image_o3d = o3d.geometry.Image(image)
+        rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(image_o3d, depth_o3d, 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)
+        return points, colors
+    
+    def generate_fig(self, image):
+        points, colors = self.generate_pcl(image)
+        data = {'x': points[:, 0], 'y': points[:, 1], 'z': points[:, 2],
+            'red': colors[:, 0], 'green': colors[:, 1], 'blue': colors[:, 2]}
+        df = pd.DataFrame(data)
+        size = np.zeros(len(df))
+        size[:] = 0.01
+        # Step 6: Create a 3D scatter plot using Plotly Express
+        fig = px.scatter_3d(df, x='x', y='y', z='z', color='red', size=size)
+        return fig
+
+    
 
 
 

utils.py

@@ -93,7 +93,9 @@ def create_3d_pc(rgb_image, depth_image, depth=10):
     return filename # Return the file path where the PLY file is saved
 
 
-def point_cloud(rgb_image, depth_image):
+def point_cloud(rgb_image):
+    depth_predictor = DepthPredictor()
+    depth_result = depth_predictor.predict(rgb_image)
     # Step 2: Create an RGBD image from the RGB and depth image
     depth_o3d = o3d.geometry.Image(depth_image)
     image_o3d = o3d.geometry.Image(rgb_image)
@@ -112,6 +114,7 @@ def point_cloud(rgb_image, depth_image):
     # Step 6: Create a 3D scatter plot using Plotly Express
     fig = px.scatter_3d(df, x='x', y='y', z='z', color='red', size=size)
     
+    
     return fig
 
 def array_PCL(rgb_image, depth_image):