Update app.py

app.py

@@ -79,18 +79,39 @@ with gr.Blocks(css=css) as demo:
 
     cmap = matplotlib.colormaps.get_cmap('Spectral_r')
 
-    def on_submit(image):
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
+def on_submit(image):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
-        original_image = image.copy()
+    original_image = image.copy()
 
-        h, w = image.size  # For PIL images, use .size instead of .shape
+    h, w = image.size  # For PIL images, use .size instead of .shape
 
-        depth = predict_depth(image)
+    depth = predict_depth(image)
 
-        raw_depth = Image.fromarray(depth)
-        tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+    # Debugging info
+    print("Type of depth:", type(depth))
+    print("Shape of depth:", depth.shape if isinstance(depth, np.ndarray) else "N/A")
+
+    if isinstance(depth, np.ndarray):
+        print("Data type of depth:", depth.dtype)
+
+    # Ensure depth is a NumPy array and has the correct shape and dtype
+    if not isinstance(depth, np.ndarray):
+        raise TypeError("Expected a NumPy array for depth, but got {}".format(type(depth)))
+
+    # Adjust the depth array if needed
+    if depth.ndim == 2:
+        # 2D array: expected for grayscale depth maps
+        depth = depth.astype('uint16')  # Convert to a suitable type
+    elif depth.ndim == 3 and depth.shape[2] == 1:
+        # 3D array with a single channel (e.g., shape (H, W, 1))
+        depth = depth[:, :, 0].astype('uint16')
+    else:
+        raise ValueError("Unsupported depth array shape: {}".format(depth.shape))
+
+    # Now convert to a PIL Image
+    raw_depth = Image.fromarray(depth)        tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
         raw_depth.save(tmp_raw_depth.name)
 
         depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0