depth-pro / app.py
akhaliq's picture
akhaliq HF staff
Update app.py
2e549d0 verified
raw
history blame
2.21 kB
import gradio as gr
from PIL import Image
import src.depth_pro as depth_pro
import numpy as np
import matplotlib.pyplot as plt
import subprocess
import spaces
import torch
# Run the script to get pretrained models
subprocess.run(["bash", "get_pretrained_models.sh"])
# Load model and preprocessing transform
model, transform = depth_pro.create_model_and_transforms()
model.eval()
@spaces.GPU(duration=120)
def predict_depth(input_image):
try:
# Preprocess the image
result = depth_pro.load_rgb(input_image)
image = result[0]
f_px = result[-1] # Assuming f_px is the last item in the returned tuple
image = transform(image)
# Run inference
prediction = model.infer(image, f_px=f_px)
depth = prediction["depth"] # Depth in [m]
focallength_px = prediction["focallength_px"] # Focal length in pixels
# Convert depth to numpy array if it's a torch tensor
if isinstance(depth, torch.Tensor):
depth = depth.cpu().numpy()
# Ensure depth is a 2D numpy array
if depth.ndim != 2:
depth = depth.squeeze()
# Normalize depth for visualization
depth_min = np.min(depth)
depth_max = np.max(depth)
depth_normalized = (depth - depth_min) / (depth_max - depth_min)
# Create a color map
plt.figure(figsize=(10, 10))
plt.imshow(depth_normalized, cmap='viridis')
plt.colorbar(label='Depth')
plt.title('Predicted Depth Map')
plt.axis('off')
# Save the plot to a file
output_path = "depth_map.png"
plt.savefig(output_path)
plt.close()
return output_path, f"Focal length: {focallength_px:.2f} pixels"
except Exception as e:
return None, f"An error occurred: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=predict_depth,
inputs=gr.Image(type="filepath"),
outputs=[gr.Image(type="filepath", label="Depth Map"), gr.Textbox(label="Focal Length or Error Message")],
title="Depth Prediction Demo",
description="Upload an image to predict its depth map and focal length."
)
# Launch the interface
iface.launch()