Alessio Grancini
commited on
Update app.py
Browse files
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 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
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 |
-
|
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 |
-
|
229 |
-
|
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
|