Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import json | |
| from unidepth.models import UniDepthV2 | |
| import os | |
| import matplotlib.pyplot as plt | |
| import matplotlib | |
| from PIL import Image | |
| # Load model configurations and initialize model | |
| def load_model(config_path, model_path, encoder, device): | |
| with open(config_path) as f: | |
| config = json.load(f) | |
| model = UniDepthV2(config) | |
| model.load_state_dict(torch.load(model_path, map_location=device)['model'], strict=True) | |
| model = model.to(device).eval() | |
| return model | |
| # Inference function | |
| def depth_estimation(image, model_path, encoder='vits'): | |
| try: | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| config_path = 'configs/config_v2_vits14.json' | |
| # Ensure model path exists or download if needed | |
| if not os.path.exists(model_path): | |
| return "Model checkpoint not found. Please upload a valid model path." | |
| model = load_model(config_path, model_path, encoder, device) | |
| # Preprocess image | |
| rgb = torch.from_numpy(np.array(image)).permute(2, 0, 1).to(device) # C, H, W | |
| predictions = model.infer(rgb) | |
| depth = predictions["depth"].squeeze().to('cpu').numpy() | |
| min_depth = depth.min() | |
| max_depth = depth.max() | |
| depth_normalized = (depth - min_depth) / (max_depth - min_depth) | |
| # Apply colormap | |
| cmap = matplotlib.colormaps.get_cmap('Spectral') | |
| depth_color = (cmap(depth_normalized)[:, :, :3] * 255).astype(np.uint8) | |
| # Create a figure and axis for the colorbar | |
| fig, ax = plt.subplots(figsize=(6, 0.4)) | |
| fig.subplots_adjust(bottom=0.5) | |
| # Create a colorbar | |
| norm = matplotlib.colors.Normalize(vmin=min_depth, vmax=max_depth) | |
| sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) | |
| sm.set_array([]) | |
| cbar = fig.colorbar(sm, cax=ax, orientation='horizontal', label='Depth (meters)') | |
| # Save the colorbar to a BytesIO object | |
| from io import BytesIO | |
| buf = BytesIO() | |
| fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1) | |
| plt.close(fig) | |
| buf.seek(0) | |
| # Open the colorbar image | |
| colorbar_img = Image.open(buf) | |
| # Create a new image with space for the colorbar | |
| new_height = depth_color.shape[0] + colorbar_img.size[1] | |
| new_img = Image.new('RGB', (depth_color.shape[1], new_height), (255, 255, 255)) | |
| # Paste the depth image and colorbar | |
| new_img.paste(Image.fromarray(depth_color), (0, 0)) | |
| new_img.paste(colorbar_img, (0, depth_color.shape[0])) | |
| return new_img | |
| except Exception as e: | |
| return f"Error occurred: {str(e)}" | |
| # Gradio Interface | |
| def main(): | |
| iface = gr.Interface( | |
| fn=depth_estimation, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Input Image"), | |
| gr.Textbox(value='checkpoint/latest.pth', label='Model Path'), | |
| gr.Dropdown(choices=['vits', 'vitb', 'vitl', 'vitg'], value='vits', label='Encoder'), | |
| ], | |
| outputs=[ | |
| gr.Image(type="pil", label="Predicted Depth") | |
| ], | |
| title="Depth Anything V2 Metric Depth Estimation", | |
| description="Upload an image to get its estimated depth map using Depth Anything V2.", | |
| ) | |
| iface.launch() | |
| if __name__ == "__main__": | |
| main() | |