import torch import torch.nn.functional as F import logging import os import os.path as osp #os.system('nvidia-smi') import cupy import sys CODE_SPACE=os.path.dirname(os.path.dirname(os.path.abspath(__file__))) try: from mmcv.utils import Config, DictAction except: from mmengine import Config, DictAction from mono.utils.logger import setup_logger import glob from mono.utils.comm import init_env from mono.model.monodepth_model import get_configured_monodepth_model from mono.utils.running import load_ckpt from mono.utils.do_test import transform_test_data_scalecano, get_prediction from mono.utils.custom_data import load_from_annos, load_data from mono.utils.avg_meter import MetricAverageMeter from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs, save_normal_val_imgs import cv2 from tqdm import tqdm import numpy as np from PIL import Image import matplotlib.pyplot as plt from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud from mono.utils.transform import gray_to_colormap from mono.utils.visualization import vis_surface_normal import gradio as gr import fire from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud from datetime import datetime import time #torch.hub.download_url_to_file('https://images.unsplash.com/photo-1437622368342-7a3d73a34c8f', 'turtle.jpg') #torch.hub.download_url_to_file('https://images.unsplash.com/photo-1519066629447-267fffa62d4b', 'lions.jpg') cfg_large = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.large.py') model_large = get_configured_monodepth_model(cfg_large, ) model_large, _, _, _ = load_ckpt('./weight/metric_depth_vit_large_800k.pth', model_large, strict_match=False) model_large.eval() cfg_small = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.small.py') model_small = get_configured_monodepth_model(cfg_small, ) model_small, _, _, _ = load_ckpt('./weight/metric_depth_vit_small_800k.pth', model_small, strict_match=False) model_small.eval() device = "cuda" model_large.to(device) model_small.to(device) outputs_dir = "./outs" def depth_normal(img_path, model_selection="vit-small"): if model_selection == "vit-small": model = model_small cfg = cfg_small elif model_selection == "vit-large": model = model_large cfg = cfg_large else: raise NotImplementedError img = Image.open(img_path) cv_image = np.array(img) img = cv_image img = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) intrinsic = [1000.0, 1000.0, img.shape[1]/2, img.shape[0]/2] rgb_input, cam_models_stacks, pad, label_scale_factor = transform_test_data_scalecano(img, intrinsic, cfg.data_basic) with torch.no_grad(): pred_depth, pred_depth_scale, scale, output = get_prediction( model = model, input = rgb_input, cam_model = cam_models_stacks, pad_info = pad, scale_info = label_scale_factor, gt_depth = None, normalize_scale = cfg.data_basic.depth_range[1], ori_shape=[img.shape[0], img.shape[1]], ) pred_normal = output['normal_out_list'][0][:, :3, :, :] H, W = pred_normal.shape[2:] pred_normal = pred_normal[:, :, pad[0]:H-pad[1], pad[2]:W-pad[3]] pred_depth = pred_depth.squeeze().cpu().numpy() pred_depth[pred_depth<0] = 0 pred_color = gray_to_colormap(pred_depth) ##formatted = (output * 255 / np.max(output)).astype('uint8') path_output_dir = os.path.splitext(os.path.basename(img_path))[0] + datetime.now().strftime('%Y%m%d-%H%M%S') path_output_dir = os.path.join(path_output_dir, outputs_dir) os.makedirs(path_output_dir, exist_ok=True) name_base = os.path.splitext(os.path.basename(img_path))[0] depth_np = pred_depth normal_np = torch.nn.functional.interpolate(pred_normal, [img.shape[0], img.shape[1]], mode='bilinear').squeeze().cpu().numpy() normal_np = normal_np.transpose(1,2,0) pred_normal = pred_normal.squeeze() if pred_normal.size(0) == 3: pred_normal = pred_normal.permute(1,2,0) pred_color_normal = vis_surface_normal(pred_normal) depth_path = os.path.join(path_output_dir, f"{name_base}_depth.npy") normal_path = os.path.join(path_output_dir, f"{name_base}_normal.npy") np.save(normal_path, normal_np) np.save(depth_path, depth_np) ori_w = img.shape[1] ori_h = img.shape[0] img = Image.fromarray(pred_color) #img = img.resize((int(300 * ori_w/ ori_h), 300)) img_normal = Image.fromarray(pred_color_normal) #img_normal = img_normal.resize((int(300 * ori_w/ ori_h), 300)) return img, img_normal, [depth_path, normal_path] def reconstruction(img_path, files, focal_length, reconstructed_file): img = Image.open(img_path) cv_image = np.array(img) img = cv_image depth_np = np.load(files[0]) pcd = reconstruct_pcd(depth_np * focal_length / 1000, focal_length, focal_length, img.shape[1]/2, img.shape[0]/2) pcd_path = files[0].replace('_depth.npy', '.ply') save_point_cloud(pcd.reshape((-1, 3)), img.reshape(-1, 3), pcd_path) return [pcd_path] title = "Metric3D" description = "Gradio demo for Metric3D which takes in a single image for computing metric depth and surface normal. To use it, simply upload your image, or click one of the examples to load them. Learn more from our paper linked below." article = "

Metric3D: Towards Zero-shot Metric 3D Prediction from A Single Image | Github Repo

" examples = [ ["files/museum.jpg"], ["files/terra.jpg"], ["files/underwater.jpg"], ["files/venue.jpg"] ] def run_demo(): _TITLE = '''Metric3Dv2: A versatile monocular geometric foundation model for zero-shot metric depth and surface normal estimation''' _DESCRIPTION = description with gr.Blocks(title=_TITLE) as demo: with gr.Row(): with gr.Column(scale=1): gr.Markdown('# ' + _TITLE) gr.Markdown(_DESCRIPTION) with gr.Row(variant='panel'): with gr.Column(scale=1): #input_image = gr.Image(type='pil', label='Original Image') input_image = gr.Image(type='filepath', height=300, label='Input image') example_folder = os.path.join(os.path.dirname(__file__), "./files") example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)] gr.Examples( examples=example_fns, inputs=[input_image], cache_examples=False, label='Examples (click one of the images below to start)', examples_per_page=30 ) model_choice = gr.Dropdown(["vit-small", "vit-large"], label="Model", info="Select a model type", value="vit-small") run_btn = gr.Button('Predict', variant='primary', interactive=True) with gr.Column(scale=1): depth = gr.Image(interactive=False, label="Depth") normal = gr.Image(interactive=False, label="Normal") with gr.Row(): files = gr.Files( label = "Depth and Normal (numpy)", elem_id = "download", interactive=False, ) with gr.Row(): recon_btn = gr.Button('Focal Length Available? If Yes, Enter and Click Here for Metric 3D Reconstruction', variant='primary', interactive=True) focal_length = gr.Number(value=1000, label="Focal Length") with gr.Row(): reconstructed_file = gr.Files( label = "3D pointclouds (plyfile)", elem_id = "download", interactive=False ) run_btn.click(fn=depth_normal, inputs=[input_image, model_choice], outputs=[depth, normal, files] ) recon_btn.click(fn=reconstruction, inputs=[input_image, files, focal_length], outputs=[reconstructed_file] ) demo.queue().launch(share=True, max_threads=80) if __name__ == '__main__': fire.Fire(run_demo)