Update app.py

app.py

@@ -107,6 +107,115 @@ def model_selector(model_type):
     img_seg = ImageSegmenter(model_type=yolo_model)
     depth_estimator = MonocularDepthEstimator(model_type=midas_model)
 
+    # START 
+    # added for lens studio 
+
+    
+def get_box_vertices(bbox):
+    """Convert bbox to corner vertices"""
+    x1, y1, x2, y2 = bbox
+    return [
+        [x1, y1],  # top-left
+        [x2, y1],  # top-right
+        [x2, y2],  # bottom-right
+        [x1, y2]   # bottom-left
+    ]
+
+def depth_at_center(depth_map, bbox):
+    """Get depth at center of bounding box"""
+    x1, y1, x2, y2 = bbox
+    center_x = int((x1 + x2) / 2)
+    center_y = int((y1 + y2) / 2)
+    
+    # Sample a small region around center for stability
+    region = depth_map[
+        max(0, center_y-2):min(depth_map.shape[0], center_y+3),
+        max(0, center_x-2):min(depth_map.shape[1], center_x+3)
+    ]
+    return np.median(region)
+
+def get_camera_matrix(depth_estimator):
+    """Get camera calibration matrix"""
+    return {
+        "fx": depth_estimator.fx_depth,
+        "fy": depth_estimator.fy_depth,
+        "cx": depth_estimator.cx_depth,
+        "cy": depth_estimator.cy_depth
+    }
+
+@spaces.GPU
+def get_detection_data(image):
+    """
+    Process image and return structured detection data with camera parameters
+    """
+    try:
+        # Resize image if needed
+        image = utils.resize(image)
+        
+        # Run detections
+        image_segmentation, objects_data = img_seg.predict(image)
+        depthmap, depth_colormap = depth_estimator.make_prediction(image)
+        
+        # Get original image dimensions
+        height, width = image.shape[:2]
+        
+        # Scale factor for normalizing coordinates
+        scale_x = width / depthmap.shape[1]
+        scale_y = height / depthmap.shape[0]
+        
+        # Process each detection
+        detections = []
+        for obj in objects_data:
+            cls_id, category, center, mask, color = obj
+            
+            # Get bounding box (assuming it's available in objects_data)
+            bbox = get_object_bbox(mask)  # You'll need to implement this
+            
+            # Get normalized coordinates
+            bbox_norm = [
+                bbox[0] / width,
+                bbox[1] / height,
+                bbox[2] / width,
+                bbox[3] / height
+            ]
+            
+            # Get vertices
+            vertices = get_box_vertices(bbox_norm)
+            
+            # Get depth
+            depth_value = depth_at_center(depthmap, bbox)
+            
+            # Create detection object
+            detection = {
+                "category": category,
+                "confidence": 1.0,  # Add actual confidence if available
+                "bbox": bbox_norm,
+                "depth": float(depth_value),  # Convert to native Python float
+                "vertices": vertices,
+                "color": [float(c/255) for c in color],  # Normalize color
+                "mask": mask.tolist() if isinstance(mask, np.ndarray) else mask
+            }
+            detections.append(detection)
+        
+        # Prepare response
+        response = {
+            "detections": detections,
+            "depth_map": depthmap.tolist(),
+            "camera_params": get_camera_matrix(depth_estimator),
+            "image_size": {
+                "width": width,
+                "height": height
+            }
+        }
+        
+        return response
+        
+    except Exception as e:
+        print(f"Error in get_detection_data: {str(e)}")
+        raise
+
+    # ENDS
+
 def cancel():
     CANCEL_PROCESSING = True