jens commited on
Commit
7544b79
·
1 Parent(s): 3756abb

map depth attempt

Browse files
Files changed (2) hide show
  1. app.py +1 -1
  2. inference.py +32 -9
app.py CHANGED
@@ -159,7 +159,7 @@ with block:
159
 
160
  def on_depth_reconstruction_btn_click(inputs):
161
  print("depth reconstruction")
162
- path = dpt.generate_obj_masks2(image=inputs[input_image],
163
  cube_size=inputs[cube_size],
164
  n_samples=inputs[n_samples],
165
  masks=inputs[masks],
 
159
 
160
  def on_depth_reconstruction_btn_click(inputs):
161
  print("depth reconstruction")
162
+ path = dpt.generate_obj_rgb(image=inputs[input_image],
163
  cube_size=inputs[cube_size],
164
  n_samples=inputs[n_samples],
165
  masks=inputs[masks],
inference.py CHANGED
@@ -15,7 +15,7 @@ import matplotlib.pyplot as plt
15
 
16
 
17
 
18
- def map_image_range(image, min_value, max_value):
19
  """
20
  Maps the values of a numpy image array to a specified range.
21
 
@@ -28,11 +28,19 @@ def map_image_range(image, min_value, max_value):
28
  numpy.ndarray: Image array with values mapped to the specified range.
29
  """
30
  # Ensure the input image is a numpy array
31
- print(np.min(image))
32
- print(np.max(image))
33
-
 
 
 
 
 
 
34
  # Map the values to the specified range
35
- mapped_image = (image - 0) * (max_value - min_value) / (1 - 0) + min_value
 
 
36
  return mapped_image
37
 
38
  def PCL(mask, depth):
@@ -63,6 +71,20 @@ def PCL(mask, depth):
63
  colors = colors[mask]
64
  return points, colors
65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66
  class DepthPredictor:
67
  def __init__(self):
68
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -124,11 +146,12 @@ class DepthPredictor:
124
  ax.scatter(points,size=0.01, c=colors, marker='o')
125
  return fig
126
 
127
- def generate_obj_rgb(self, image, n_samples, cube_size):
128
  # Step 1: Create a point cloud
129
- point_cloud, color_array = self.generate_pcl(image)
130
- #point_cloud, color_array = point_cloud[mask.ravel()[:-1]], color_array[mask.ravel()[:-1]]
131
- # sample 1000 points
 
132
  idxs = np.random.choice(len(point_cloud), int(n_samples))
133
  point_cloud = point_cloud[idxs]
134
  color_array = color_array[idxs]
 
15
 
16
 
17
 
18
+ def map_image_range(depth, min_value, max_value):
19
  """
20
  Maps the values of a numpy image array to a specified range.
21
 
 
28
  numpy.ndarray: Image array with values mapped to the specified range.
29
  """
30
  # Ensure the input image is a numpy array
31
+ print(np.min(depth))
32
+ print(np.max(depth))
33
+ depth = np.array(depth)
34
+ # map the depth values are between 0 and 1
35
+ depth = (depth - depth.min()) / (depth.max() - depth.min())
36
+ # invert
37
+ depth = 1 - depth
38
+ print(np.min(depth))
39
+ print(np.max(depth))
40
  # Map the values to the specified range
41
+ mapped_image = (depth - 0) * (max_value - min_value) / (1 - 0) + min_value
42
+ print(np.min(mapped_image))
43
+ print(np.max(mapped_image))
44
  return mapped_image
45
 
46
  def PCL(mask, depth):
 
71
  colors = colors[mask]
72
  return points, colors
73
 
74
+ def PCL_rgb(rgb, depth):
75
+ assert rgb.shape == depth.shape
76
+ assert type(rgb) == np.ndarray
77
+ assert type(depth) == np.ndarray
78
+ depth_o3d = o3d.geometry.Image(depth)
79
+ image_o3d = o3d.geometry.Image(rgb)
80
+ rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(image_o3d, depth_o3d, convert_rgb_to_intensity=False)
81
+ # Step 3: Create a PointCloud from the RGBD image
82
+ pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault))
83
+ # Step 4: Convert PointCloud data to a NumPy array
84
+ points = np.asarray(pcd.points)
85
+ colors = np.asarray(pcd.colors)
86
+ return points, colors
87
+
88
  class DepthPredictor:
89
  def __init__(self):
90
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
146
  ax.scatter(points,size=0.01, c=colors, marker='o')
147
  return fig
148
 
149
+ def generate_obj_rgb(self, image, n_samples, cube_size, max_depth, min_depth):
150
  # Step 1: Create a point cloud
151
+ depth = self.predict(image)
152
+ depth = map_image_range(depth, min_depth, max_depth)
153
+ point_cloud, color_array = PCL_rgb(image, depth)
154
+ image = np.array(image)
155
  idxs = np.random.choice(len(point_cloud), int(n_samples))
156
  point_cloud = point_cloud[idxs]
157
  color_array = color_array[idxs]