jens commited on
Commit
e6c4109
·
1 Parent(s): c09fa67

boolean dimension fix

Browse files
Files changed (1) hide show
  1. inference.py +5 -1
inference.py CHANGED
@@ -98,7 +98,11 @@ class DepthPredictor:
98
  # Create cubes and add them to the mesh
99
  cs = [(255,0,0),(0,255,0),(0,0,255)]
100
  for c,(mask, _) in zip(cs, masks):
101
- mask = mask.ravel()[:-1]
 
 
 
 
102
  point_cloud_subset, color_array_subset = point_cloud[mask], color_array[mask]
103
  idxs = np.random.choice(len(point_cloud_subset), int(n_samples))
104
  point_cloud_subset = point_cloud_subset[idxs]
 
98
  # Create cubes and add them to the mesh
99
  cs = [(255,0,0),(0,255,0),(0,0,255)]
100
  for c,(mask, _) in zip(cs, masks):
101
+ if len(mask) == len(point_cloud):
102
+ mask = mask.ravel()
103
+ else:
104
+ mask = mask.ravel()[:-1]
105
+
106
  point_cloud_subset, color_array_subset = point_cloud[mask], color_array[mask]
107
  idxs = np.random.choice(len(point_cloud_subset), int(n_samples))
108
  point_cloud_subset = point_cloud_subset[idxs]