Alessio Grancini commited on
Commit
a4031b7
·
verified ·
1 Parent(s): 4512499

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -38
app.py CHANGED
@@ -146,22 +146,28 @@ def get_camera_matrix(depth_estimator):
146
  @spaces.GPU
147
  def get_detection_data(image):
148
  """Get structured detection data with depth information"""
 
 
149
  try:
150
  # Resize image to standard size
151
  image = utils.resize(image)
152
-
 
 
 
 
153
  # Get detections and depth
154
  image_segmentation, objects_data = img_seg.predict(image)
155
  depthmap, depth_colormap = depth_estimator.make_prediction(image)
156
-
157
  # Process each detection
158
  detections = []
159
  for data in objects_data:
160
  cls_id, cls_name, cls_center, cls_mask, cls_clr = data
161
-
162
  # Get masked depth for this object
163
  masked_depth, mean_depth = utils.get_masked_depth(depthmap, cls_mask)
164
-
165
  # Get bounding box from mask
166
  y_indices, x_indices = np.where(cls_mask > 0)
167
  if len(x_indices) > 0 and len(y_indices) > 0:
@@ -169,64 +175,61 @@ def get_detection_data(image):
169
  y1, y2 = np.min(y_indices), np.max(y_indices)
170
  else:
171
  continue
172
-
173
  # Normalize coordinates
174
- height, width = image.shape[:2]
175
  bbox_normalized = [
176
- float(x1/width),
177
- float(y1/height),
178
- float(x2/width),
179
- float(y2/height)
180
  ]
181
-
182
  detection = {
183
  "id": int(cls_id),
184
  "category": cls_name,
185
  "center": [
186
- float(cls_center[0]/width),
187
- float(cls_center[1]/height)
188
  ],
189
  "bbox": bbox_normalized,
190
  "depth": float(mean_depth * 10), # Convert to meters as done in utils
191
- "color": [float(c/255) for c in cls_clr],
192
  "mask": cls_mask.tolist(),
193
- "confidence": 1.0 # Add actual confidence if available
194
  }
195
  detections.append(detection)
196
-
197
- # Get camera parameters from depth estimator
198
- camera_params = {
199
- "fx": depth_estimator.fx_depth,
200
- "fy": depth_estimator.fy_depth,
201
- "cx": depth_estimator.cx_depth,
202
- "cy": depth_estimator.cy_depth
203
- }
204
-
 
 
 
 
205
  # Generate point cloud data if needed
206
  point_clouds = utils.generate_obj_pcd(depthmap, objects_data)
207
  pcd_data = [
208
- {
209
- "points": np.asarray(pcd.points).tolist(),
210
- "color": [float(c/255) for c in color]
211
- }
212
  for pcd, color in point_clouds
213
  ]
214
-
215
  return {
216
  "detections": detections,
217
  "depth_map": depthmap.tolist(),
218
  "camera_params": camera_params,
219
- "image_size": {
220
- "width": width,
221
- "height": height
222
- },
223
- "point_clouds": pcd_data
224
  }
225
-
226
  except Exception as e:
227
- print(f"Error in get_detection_data: {str(e)}")
228
- raise
229
- # ENDS
230
 
231
  def cancel():
232
  CANCEL_PROCESSING = True
 
146
  @spaces.GPU
147
  def get_detection_data(image):
148
  """Get structured detection data with depth information"""
149
+ width, height = 640, 480 # Set default values to avoid UnboundLocalError
150
+
151
  try:
152
  # Resize image to standard size
153
  image = utils.resize(image)
154
+
155
+ # Ensure width and height are properly set
156
+ if hasattr(image, "shape"):
157
+ height, width = image.shape[:2] # Extract actual dimensions
158
+
159
  # Get detections and depth
160
  image_segmentation, objects_data = img_seg.predict(image)
161
  depthmap, depth_colormap = depth_estimator.make_prediction(image)
162
+
163
  # Process each detection
164
  detections = []
165
  for data in objects_data:
166
  cls_id, cls_name, cls_center, cls_mask, cls_clr = data
167
+
168
  # Get masked depth for this object
169
  masked_depth, mean_depth = utils.get_masked_depth(depthmap, cls_mask)
170
+
171
  # Get bounding box from mask
172
  y_indices, x_indices = np.where(cls_mask > 0)
173
  if len(x_indices) > 0 and len(y_indices) > 0:
 
175
  y1, y2 = np.min(y_indices), np.max(y_indices)
176
  else:
177
  continue
178
+
179
  # Normalize coordinates
 
180
  bbox_normalized = [
181
+ float(x1 / width),
182
+ float(y1 / height),
183
+ float(x2 / width),
184
+ float(y2 / height),
185
  ]
186
+
187
  detection = {
188
  "id": int(cls_id),
189
  "category": cls_name,
190
  "center": [
191
+ float(cls_center[0] / width),
192
+ float(cls_center[1] / height),
193
  ],
194
  "bbox": bbox_normalized,
195
  "depth": float(mean_depth * 10), # Convert to meters as done in utils
196
+ "color": [float(c / 255) for c in cls_clr],
197
  "mask": cls_mask.tolist(),
198
+ "confidence": 1.0, # Add actual confidence if available
199
  }
200
  detections.append(detection)
201
+
202
+ # Get camera parameters from depth estimator (check if attributes exist)
203
+ try:
204
+ camera_params = {
205
+ "fx": getattr(depth_estimator, "fx_depth", 0),
206
+ "fy": getattr(depth_estimator, "fy_depth", 0),
207
+ "cx": getattr(depth_estimator, "cx_depth", width // 2),
208
+ "cy": getattr(depth_estimator, "cy_depth", height // 2),
209
+ }
210
+ except AttributeError:
211
+ print("⚠️ Camera parameters are not properly set in depth_estimator.")
212
+ camera_params = {"fx": 0, "fy": 0, "cx": width // 2, "cy": height // 2}
213
+
214
  # Generate point cloud data if needed
215
  point_clouds = utils.generate_obj_pcd(depthmap, objects_data)
216
  pcd_data = [
217
+ {"points": np.asarray(pcd.points).tolist(), "color": [float(c / 255) for c in color]}
 
 
 
218
  for pcd, color in point_clouds
219
  ]
220
+
221
  return {
222
  "detections": detections,
223
  "depth_map": depthmap.tolist(),
224
  "camera_params": camera_params,
225
+ "image_size": {"width": width, "height": height},
226
+ "point_clouds": pcd_data,
 
 
 
227
  }
228
+
229
  except Exception as e:
230
+ print(f"🚨 Error in get_detection_data: {str(e)}")
231
+ return {"error": str(e)}
232
+
233
 
234
  def cancel():
235
  CANCEL_PROCESSING = True