Update app.py

app.py

@@ -7,7 +7,7 @@ from typing import Union
 from pathlib import Path
 import os
 
-def predict_depth(image: Image.Image, auto_rotate: bool, remove_alpha: bool, model, transform):
+def predict_depth(image: Image.Image, auto_rotate: bool, remove_alpha: bool, grayscale: bool, model, transform):
     # Convert the PIL image to a temporary file path if needed
     image_path = "temp_image.jpg"
     image.save(image_path)
@@ -30,14 +30,20 @@ def predict_depth(image: Image.Image, auto_rotate: bool, remove_alpha: bool, mod
 
     focallength = prediction["focallength_px"].cpu().numpy()
 
-    # Normalize and colorize depth map
-    cmap = plt.get_cmap("turbo_r")
-    color_depth = (cmap(inverse_depth_normalized)[..., :3] * 255).astype(np.uint8)
+    if grayscale:
+        # Normalize the inverse depth map to 0-255 and convert to grayscale
+        grayscale_depth = (inverse_depth_normalized * 255).astype(np.uint8)
+        depth_image = Image.fromarray(grayscale_depth, mode="L")
+    else:
+        # Normalize and colorize depth map
+        cmap = plt.get_cmap("turbo_r")
+        color_depth = (cmap(inverse_depth_normalized)[..., :3] * 255).astype(np.uint8)
+        depth_image = Image.fromarray(color_depth)
 
     # Clean up temporary image
     os.remove(image_path)
 
-    return Image.fromarray(color_depth), focallength  # Return depth map and f_px
+    return depth_image, focallength  # Return depth map and f_px
 
 def main():
     # Load model and preprocessing transform
@@ -46,11 +52,12 @@ def main():
 
     # Set up Gradio interface
     iface = gr.Interface(
-        fn=lambda image, auto_rotate, remove_alpha: predict_depth(image, auto_rotate, remove_alpha, model, transform),
+        fn=lambda image, auto_rotate, remove_alpha, grayscale: predict_depth(image, auto_rotate, remove_alpha, grayscale, model, transform),
         inputs=[
             gr.Image(type="pil", label="Upload Image"),  # Use image browser for input
             gr.Checkbox(label="Auto Rotate", value=True),  # Checkbox for auto_rotate
-            gr.Checkbox(label="Remove Alpha", value=True)   # Checkbox for remove_alpha
+            gr.Checkbox(label="Remove Alpha", value=True),  # Checkbox for remove_alpha
+            gr.Checkbox(label="Grayscale Depth", value=False)  # Checkbox for grayscale
         ],
         outputs=[
             gr.Image(label="Depth Map"),  # Use PIL image output
@@ -65,4 +72,4 @@ def main():
     iface.launch()
 
 if __name__ == "__main__":
-    main()
+    main()