Update app.py

app.py

@@ -1,6 +1,11 @@
 import gradio as gr
 from PIL import Image
 import numpy as np
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+
+# Load the Marigold model
+feature_extractor = AutoFeatureExtractor.from_pretrained("prs-eth/marigold-depth-lcm-v1-0")
+model = AutoModelForImageClassification.from_pretrained("prs-eth/marigold-depth-lcm-v1-0")
 
 def load_image(image):
     """ Convert uploaded image to grayscale. """
@@ -33,23 +38,30 @@ def create_normal_map(image):
 
     return Image.fromarray(normal_map, 'RGB')
 
-def convert_to_grayscale(image):
-    """ Convert the normal map to a grayscale image. """
-    grayscale_image = image.convert('L')
-    return grayscale_image
+def estimate_depth(image):
+    """ Estimate depth using the Marigold model. """
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    depth_map = outputs.logits  # Adjust if the output tensor is named differently
+    return depth_map
 
 def interface(image):
     normal_map = create_normal_map(image)
-    grayscale_image = convert_to_grayscale(normal_map)
-    return normal_map, grayscale_image
+    grayscale_image = load_image(normal_map)
+    depth_map = estimate_depth(image)
+    return normal_map, grayscale_image, depth_map
 
 # Set up the Gradio interface
 iface = gr.Interface(
     fn=interface,
     inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=[gr.Image(type="pil", label="Normal Map"), gr.Image(type="pil", label="Grayscale Image")],
-    title="Normal Map and Grayscale Generator",
-    description="Upload an image to generate its normal map and a grayscale version of the normal map. Both images retain the original dimensions."
+    outputs=[
+        gr.Image(type="pil", label="Normal Map"),
+        gr.Image(type="pil", label="Grayscale Image"),
+        gr.Image(type="pil", label="Depth Map")  # Adjust the output type if needed
+    ],
+    title="Normal Map, Grayscale, and Depth Map Generator",
+    description="Upload an image to generate its normal map, a grayscale version, and estimate its depth."
 )
 
 iface.launch()