Update app.py

app.py

@@ -1,47 +1,6 @@
 # Importing the requirements
-import numpy as np
 import gradio as gr
-import torch
-from PIL import Image
-from transformers import DPTImageProcessor, DPTForDepthEstimation
-
-
-# Load the model and feature extractor
-feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-beit-large-512")
-model = DPTForDepthEstimation.from_pretrained("Intel/dpt-beit-large-512")
-
-
-def process_image(image):
-    """
-    Preprocesses an image, passes it through a model, and returns the formatted depth map as an image.
-
-    Args:
-        image (PIL.Image.Image): The input image.
-
-    Returns:
-        PIL.Image.Image: The formatted depth map as an image.
-    """
-
-    # Preprocess the image for the model
-    encoding = feature_extractor(image, return_tensors="pt")
-
-    # Forward pass through the model
-    with torch.no_grad():
-        outputs = model(**encoding)
-        predicted_depth = outputs.predicted_depth
-
-    # Interpolate the predicted depth map to the original image size
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=image.size[::-1],
-        mode="bicubic",
-        align_corners=False,
-    ).squeeze()
-    output = prediction.cpu().numpy()
-    formatted = (output * 255 / np.max(output)).astype("uint8")
-
-    # Return the formatted depth map as an image
-    return Image.fromarray(formatted)
+from depth_estimation import process_image  # Import the depth estimation function
 
 
 # Image input for the interface
@@ -52,9 +11,9 @@ answer = gr.Image(type="pil", label="Depth Map")
 
 # Examples for the interface
 examples = [
-    ["cat.jpg"],
-    ["dog.jpg"],
-    ["bird.jpg"],
+    ["images/cat.jpg"],
+    ["images/dog.jpg"],
+    ["images/bird.jpg"],
 ]
 
 # Title, description, and article for the interface