Update app.py

app.py

@@ -7,6 +7,9 @@ import os
 import torch
 import utils
 import plotly.graph_objects as go
+from io import BytesIO
+from PIL import Image
+import base64
 
 from image_segmenter import ImageSegmenter
 from monocular_depth_estimator import MonocularDepthEstimator
@@ -151,6 +154,7 @@ def get_detection_data(image):
     def decode_base64_image(base64_string):
         """Decodes Base64 string into a NumPy image."""
         try:
+            print(f"🔍 Received Base64 String (Truncated): {base64_string[:50]}...")  # Debugging
             img_data = base64.b64decode(base64_string)
             img = Image.open(BytesIO(img_data))
             img = np.array(img)
@@ -161,23 +165,27 @@ def get_detection_data(image):
 
     def encode_base64_image(image):
         """Encodes a NumPy image into a Base64 string."""
-        _, buffer = cv2.imencode('.png', image)
-        return base64.b64encode(buffer).decode("utf-8")
-
-    width, height = 640, 480  # Default values
+        try:
+            _, buffer = cv2.imencode('.png', image)
+            return base64.b64encode(buffer).decode("utf-8")
+        except Exception as e:
+            print(f"🚨 Error encoding image to Base64: {e}")
+            return None
 
     try:
-        if isinstance(image, str):  # Ensure we're handling a Base64 string
-            image = decode_base64_image(image)
-            if image is None:
-                return {"error": "Invalid base64 image data"}
+        if not isinstance(image, str):
+            print("🚨 Error: Expected Base64 string but received:", type(image))
+            return {"error": "Invalid input format. Expected Base64-encoded image."}
+
+        image = decode_base64_image(image)
+        if image is None:
+            return {"error": "Base64 decoding failed. Ensure correct encoding."}
 
         # Resize image
         image = utils.resize(image)
 
         # Extract dimensions
-        if hasattr(image, "shape"):
-            height, width = image.shape[:2]
+        height, width = image.shape[:2]
 
         # Get detections and depth
         image_segmentation, objects_data = img_seg.predict(image)
@@ -187,55 +195,13 @@ def get_detection_data(image):
         segmentation_b64 = encode_base64_image(image_segmentation)
         depth_b64 = encode_base64_image(depth_colormap)
 
-        # Process detections
-        detections = []
-        for data in objects_data:
-            cls_id, cls_name, cls_center, cls_mask, cls_clr = data
-            masked_depth, mean_depth = utils.get_masked_depth(depthmap, cls_mask)
-
-            y_indices, x_indices = np.where(cls_mask > 0)
-            if len(x_indices) > 0 and len(y_indices) > 0:
-                x1, x2 = np.min(x_indices), np.max(x_indices)
-                y1, y2 = np.min(y_indices), np.max(y_indices)
-            else:
-                continue
-
-            # Normalize coordinates
-            bbox_normalized = [
-                float(x1 / width),
-                float(y1 / height),
-                float(x2 / width),
-                float(y2 / height),
-            ]
-
-            detection = {
-                "id": int(cls_id),
-                "category": cls_name,
-                "center": [
-                    float(cls_center[0] / width),
-                    float(cls_center[1] / height),
-                ],
-                "bbox": bbox_normalized,
-                "depth": float(mean_depth * 10),  # Convert to meters
-                "color": [float(c / 255) for c in cls_clr],
-                "mask": cls_mask.tolist(),
-                "confidence": 1.0,  # Placeholder confidence
-            }
-            detections.append(detection)
-
-        # Camera parameters
-        camera_params = {
-            "fx": getattr(depth_estimator, "fx_depth", 0),
-            "fy": getattr(depth_estimator, "fy_depth", 0),
-            "cx": getattr(depth_estimator, "cx_depth", width // 2),
-            "cy": getattr(depth_estimator, "cy_depth", height // 2),
-        }
+        if segmentation_b64 is None or depth_b64 is None:
+            return {"error": "Failed to encode output images."}
 
         return {
-            "detections": detections,
-            "depth_map": depth_b64,  # Returning depth as Base64 image
-            "segmentation": segmentation_b64,  # Returning segmentation as Base64 image
-            "camera_params": camera_params,
+            "detections": objects_data,  # Keeping as original
+            "depth_map": depth_b64,
+            "segmentation": segmentation_b64,
             "image_size": {"width": width, "height": height},
         }