File size: 1,525 Bytes
0fc8786 f58a156 9d9399e 0fc8786 9d9399e 6f8dd42 0fc8786 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
# Importing the requirements
import numpy as np
from PIL import Image
import gradio as gr
import torch
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")
# Function to process an image and return the formatted depth map as an image
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.
"""
# Check if the image is provided
if not image:
gr.Warning("No image provided")
# 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)
|