depth mapping and n_samples

app.py

@@ -62,6 +62,8 @@ with block:
             with gr.Column():
                 pcl_figure = gr.Model3D(label="3-D Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
                 with gr.Row():
+                    max_depth = gr.Slider(minimum=0, maximum=10, step=0.01, default=1, label='Number of Samples')
+                    min_depth = gr.Slider(minimum=0, maximum=10, step=0.01, default=0.1, label='Number of Samples')
                     n_samples = gr.Slider(minimum=1e3, maximum=1e6, step=1e3, default=1e3, label='Number of Samples')
                     cube_size = gr.Slider(minimum=0.00001, maximum=0.001, step=0.000001, default=0.00001, label='Cube size')
                     depth_reconstruction_btn = gr.Button('Depth Reconstruction', variant = 'primary')
@@ -153,8 +155,12 @@ with block:
 
     def on_depth_reconstruction_btn_click(inputs):
         print("depth reconstruction")
-        image = inputs[input_image]
-        path = dpt.generate_obj_masks2(image=image, cube_size=inputs[cube_size], masks=inputs[masks]) #
+        path = dpt.generate_obj_masks2(image=inputs[input_image],
+                                      cube_size=inputs[cube_size],
+                                      n_samples=inputs[n_samples],
+                                      masks=inputs[masks],
+                                      min_depth=inputs[min_depth],
+                                      max_depth=inputs[max_depth]) #
         return {pcl_figure: path}
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)
 

inference.py

@@ -10,6 +10,7 @@ import open3d as o3d
 import pandas as pd
 import plotly.express as px
 import matplotlib.pyplot as plt
+from utils import map_image_range
 
 
 
@@ -60,9 +61,8 @@ class DepthPredictor:
                                 align_corners=False,
                         ).squeeze()
             
-        output = prediction.cpu().numpy()
-        formatted = 255 - (output * 255 / np.max(output)).astype('uint8')
-
+        output = 1 - (prediction.cpu().numpy() / np.max(output))
+        #formatted = 255 - (output * 255 / np.max(output)).astype('uint8')
         #img = Image.fromarray(formatted)
         return output
     
@@ -144,15 +144,18 @@ class DepthPredictor:
         o3d.io.write_triangle_mesh(output_file, mesh)
         return output_file
     
-    def generate_obj_masks2(self, image, masks, cube_size):
+    def generate_obj_masks2(self, image, masks, cube_size, n_samples, min_depth, max_depth):
         # Generate a point cloud
         depth = self.predict(image)
+        depth = map_image_range(depth, min_depth, max_depth)
         image = np.array(image)
         mesh = o3d.geometry.TriangleMesh()
         # Create cubes and add them to the mesh
         cs = [(255,0,0),(0,255,0),(0,0,255)]
         for c,(mask, _) in zip(cs, masks):
             points, _ = PCL(mask, depth)
+            idxs = np.random.choice(len(points), int(n_samples))
+            points = points[idxs]
             for point in points:
                 cube = o3d.geometry.TriangleMesh.create_box(width=cube_size, height=cube_size, depth=cube_size)
                 cube.translate(-point)

utils.py

@@ -195,3 +195,28 @@ def PCL3(image):
     
     return fig
 
+import numpy as np
+
+def map_image_range(image, min_value, max_value):
+    """
+    Maps the values of a numpy image array to a specified range.
+
+    Args:
+        image (numpy.ndarray): Input image array with values ranging from 0 to 1.
+        min_value (float): Minimum value of the new range.
+        max_value (float): Maximum value of the new range.
+
+    Returns:
+        numpy.ndarray: Image array with values mapped to the specified range.
+    """
+    # Ensure the input image is a numpy array
+    if not isinstance(image, np.ndarray):
+        raise ValueError("Input image must be a numpy array.")
+
+    # Ensure the image values are within the valid range (0 to 1)
+    if not (0 <= np.min(image) <= 1) or not (0 <= np.max(image) <= 1):
+        raise ValueError("Input image values must be in the range [0, 1].")
+
+    # Map the values to the specified range
+    mapped_image = (image - 0) * (max_value - min_value) / (1 - 0) + min_value
+    return mapped_image