anveshg03/Metric3D_GPU
0
1 2import torch3import torch.nn.functional as F4import logging5import os6import os.path as osp7 8#os.system('nvidia-smi')9 10import cupy11 12import sys13CODE_SPACE=os.path.dirname(os.path.dirname(os.path.abspath(__file__)))14 15try:16 from mmcv.utils import Config, DictAction17except:18 from mmengine import Config, DictAction19from mono.utils.logger import setup_logger20import glob21from mono.utils.comm import init_env22from mono.model.monodepth_model import get_configured_monodepth_model23from mono.utils.running import load_ckpt24from mono.utils.do_test import transform_test_data_scalecano, get_prediction25from mono.utils.custom_data import load_from_annos, load_data26 27from mono.utils.avg_meter import MetricAverageMeter28from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs, save_normal_val_imgs29import cv230from tqdm import tqdm31import numpy as np32from PIL import Image, ExifTags33import matplotlib.pyplot as plt34 35from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud, ply_to_obj36from mono.utils.transform import gray_to_colormap37from mono.utils.visualization import vis_surface_normal38import gradio as gr39import plotly.graph_objects as go40 41#torch.hub.download_url_to_file('https://images.unsplash.com/photo-1437622368342-7a3d73a34c8f', 'turtle.jpg')42#torch.hub.download_url_to_file('https://images.unsplash.com/photo-1519066629447-267fffa62d4b', 'lions.jpg')43 44cfg_large = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.large.py')45model_large = get_configured_monodepth_model(cfg_large, )46model_large, _, _, _ = load_ckpt('./weight/metric_depth_vit_large_800k.pth', model_large, strict_match=False)47model_large.eval()48 49cfg_small = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.small.py')50model_small = get_configured_monodepth_model(cfg_small, )51model_small, _, _, _ = load_ckpt('./weight/metric_depth_vit_small_800k.pth', model_small, strict_match=False)52model_small.eval()53 54device = "cuda"55model_large.to(device)56model_small.to(device)57 58def predict_depth_normal(img, model_selection="vit-small", fx=1000.0, fy=1000.0, state_cache={}):59 if model_selection == "vit-small":60 model = model_small61 cfg = cfg_small62 elif model_selection == "vit-large":63 model = model_large64 cfg = cfg_large65 else:66 return None, None, None, None, state_cache, "Not implemented model."67 68 if img is None:69 return None, None, None, None, state_cache, "Please upload an image and wait for the upload to complete."70 71 72 cv_image = np.array(img) 73 img = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)74 intrinsic = [fx, fy, img.shape[1]/2, img.shape[0]/2]75 rgb_input, cam_models_stacks, pad, label_scale_factor = transform_test_data_scalecano(img, intrinsic, cfg.data_basic)76 77 with torch.no_grad():78 pred_depth, pred_depth_scale, scale, output, confidence = get_prediction(79 model = model,80 input = rgb_input,81 cam_model = cam_models_stacks,82 pad_info = pad,83 scale_info = label_scale_factor,84 gt_depth = None,85 normalize_scale = cfg.data_basic.depth_range[1],86 ori_shape=[img.shape[0], img.shape[1]],87 )88 89 pred_normal = output['normal_out_list'][0][:, :3, :, :] 90 H, W = pred_normal.shape[2:]91 pred_normal = pred_normal[:, :, pad[0]:H-pad[1], pad[2]:W-pad[3]]92 93 pred_depth = pred_depth.squeeze().cpu().numpy()94 pred_depth[pred_depth<0] = 095 pred_color = gray_to_colormap(pred_depth)96 97 pred_normal = torch.nn.functional.interpolate(pred_normal, [img.shape[0], img.shape[1]], mode='bilinear').squeeze()98 pred_normal = pred_normal.permute(1,2,0)99 pred_color_normal = vis_surface_normal(pred_normal)100 pred_normal = pred_normal.cpu().numpy()101 102 # Storing depth and normal map in state for potential 3D reconstruction103 state_cache['depth'] = pred_depth104 state_cache['normal'] = pred_normal105 state_cache['img'] = img106 state_cache['intrinsic'] = intrinsic107 state_cache['confidence'] = confidence 108 109 # save depth and normal map to .npy file110 if 'save_dir' not in state_cache:111 cache_id = np.random.randint(0, 100000000000)112 while osp.exists(f'recon_cache/{cache_id:08d}'):113 cache_id = np.random.randint(0, 100000000000)114 state_cache['save_dir'] = f'recon_cache/{cache_id:08d}'115 os.makedirs(state_cache['save_dir'], exist_ok=True)116 depth_file = f"{state_cache['save_dir']}/depth.npy"117 normal_file = f"{state_cache['save_dir']}/normal.npy"118 np.save(depth_file, pred_depth)119 np.save(normal_file, pred_normal)120 121 ##formatted = (output * 255 / np.max(output)).astype('uint8')122 img = Image.fromarray(pred_color)123 img_normal = Image.fromarray(pred_color_normal)124 return img, depth_file, img_normal, normal_file, state_cache, "Success!"125 126def get_camera(img):127 if img is None:128 return None, None, None, "Please upload an image and wait for the upload to complete."129 try:130 exif = img.getexif()131 exif.update(exif.get_ifd(ExifTags.IFD.Exif))132 except:133 exif = {}134 sensor_width = exif.get(ExifTags.Base.FocalPlaneYResolution, None)135 sensor_height = exif.get(ExifTags.Base.FocalPlaneXResolution, None)136 focal_length = exif.get(ExifTags.Base.FocalLength, None)137 138 # convert sensor size to mm, see https://photo.stackexchange.com/questions/40865/how-can-i-get-the-image-sensor-dimensions-in-mm-to-get-circle-of-confusion-from139 w, h = img.size140 sensor_width = w / sensor_width * 25.4 if sensor_width is not None else None141 sensor_height = h / sensor_height * 25.4 if sensor_height is not None else None142 focal_length = focal_length * 1.0 if focal_length is not None else None143 144 message = "Success!"145 if focal_length is None:146 message = "Focal length not found in EXIF. Please manually input."147 elif sensor_width is None and sensor_height is None:148 sensor_width = 16149 sensor_height = h / w * sensor_width150 message = f"Sensor size not found in EXIF. Using {sensor_width}x{sensor_height:.2f} mm as default."151 152 return sensor_width, sensor_height, focal_length, message153 154def get_intrinsic(img, sensor_width, sensor_height, focal_length):155 if img is None:156 return None, None, "Please upload an image and wait for the upload to complete."157 if sensor_width is None or sensor_height is None or focal_length is None:158 return 1000, 1000, "Insufficient information. Try detecting camera first or use default 1000 for fx and fy."159 if sensor_width == 0 or sensor_height == 0 or focal_length == 0:160 return 1000, 1000, "Insufficient information. Try detecting camera first or use default 1000 for fx and fy."161 162 # calculate focal length in pixels163 w, h = img.size164 fx = w / sensor_width * focal_length if sensor_width is not None else None165 fy = h / sensor_height * focal_length if sensor_height is not None else None166 167 # if fx is None:168 # return fy, fy, "Sensor width not provided, using fy for both fx and fy"169 # if fy is None:170 # return fx, fx, "Sensor height not provided, using fx for both fx and fy"171 172 return fx, fy, "Success!"173 174 175def unprojection_pcd(state_cache):176 depth_map = state_cache.get('depth', None)177 normal_map = state_cache.get('normal', None)178 img = state_cache.get('img', None)179 intrinsic = state_cache.get('intrinsic', None)180 181 if depth_map is None or img is None:182 return None, "Please predict depth and normal first."183 184 # # downsample/upsample the depth map to confidence map size185 # confidence = state_cache.get('confidence', None)186 # if confidence is not None:187 # H, W = confidence.shape188 # # intrinsic[0] *= W / depth_map.shape[1]189 # # intrinsic[1] *= H / depth_map.shape[0]190 # # intrinsic[2] *= W / depth_map.shape[1]191 # # intrinsic[3] *= H / depth_map.shape[0]192 # depth_map = cv2.resize(depth_map, (W, H), interpolation=cv2.INTER_LINEAR)193 # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_LINEAR)194 195 # # filter out depth map by confidence196 # mask = confidence.cpu().numpy() > 0197 198 # downsample the depth map if too large199 if depth_map.shape[0] > 1080:200 scale = 1080 / depth_map.shape[0]201 depth_map = cv2.resize(depth_map, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)202 img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)203 intrinsic = [intrinsic[0]*scale, intrinsic[1]*scale, intrinsic[2]*scale, intrinsic[3]*scale]204 205 if 'save_dir' not in state_cache:206 cache_id = np.random.randint(0, 100000000000)207 while osp.exists(f'recon_cache/{cache_id:08d}'):208 cache_id = np.random.randint(0, 100000000000)209 state_cache['save_dir'] = f'recon_cache/{cache_id:08d}'210 os.makedirs(state_cache['save_dir'], exist_ok=True)211 212 pcd_ply = f"{state_cache['save_dir']}/output.ply"213 pcd_obj = pcd_ply.replace(".ply", ".obj")214 215 pcd = reconstruct_pcd(depth_map, intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3])216 # if mask is not None:217 # pcd_filtered = pcd[mask]218 # img_filtered = img[mask]219 pcd_filtered = pcd.reshape(-1, 3)220 img_filtered = img.reshape(-1, 3)221 222 save_point_cloud(pcd_filtered, img_filtered, pcd_ply, binary=False)223 # ply_to_obj(pcd_ply, pcd_obj)224 225 # downsample the point cloud for visualization226 num_samples = 250000227 if pcd_filtered.shape[0] > num_samples:228 indices = np.random.choice(pcd_filtered.shape[0], num_samples, replace=False)229 pcd_downsampled = pcd_filtered[indices]230 img_downsampled = img_filtered[indices]231 else:232 pcd_downsampled = pcd_filtered233 img_downsampled = img_filtered234 235 # plotly show236 color_str = np.array([f"rgb({r},{g},{b})" for b,g,r in img_downsampled])237 data=[go.Scatter3d(238 x=pcd_downsampled[:,0],239 y=pcd_downsampled[:,1],240 z=pcd_downsampled[:,2],241 mode='markers',242 marker=dict(243 size=1,244 color=color_str,245 opacity=0.8,246 )247 )]248 layout = go.Layout(249 margin=dict(l=0, r=0, b=0, t=0),250 scene=dict(251 camera = dict(252 eye=dict(x=0, y=0, z=-1),253 up=dict(x=0, y=-1, z=0)254 ),255 xaxis=dict(showgrid=False, showticklabels=False, visible=False),256 yaxis=dict(showgrid=False, showticklabels=False, visible=False),257 zaxis=dict(showgrid=False, showticklabels=False, visible=False),258 )259 )260 fig = go.Figure(data=data, layout=layout)261 262 return fig, pcd_ply, "Success!"263 264 265# title = "Metric3D"266# description = '''# Metric3Dv2: A versatile monocular geometric foundation model for zero-shot metric depth and surface normal estimation267# Gradio demo for Metric3D v1/v2 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.'''268# article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2307.10984.pdf'>Metric3D arxiv</a> | <a href='https://arxiv.org/abs/2404.15506'>Metric3Dv2 arxiv</a> | <a href='https://github.com/YvanYin/Metric3D'>Github Repo</a></p>"269 270# custom_css = '''#button1, #button2 {271# width: 20px;272# }'''273 274# examples = [275# #["turtle.jpg"],276# #["lions.jpg"]277# #["files/gundam.jpg"],278# "files/p50_pro.jpg",279# "files/iphone13.JPG",280# "files/canon_cat.JPG",281# "files/canon_dog.JPG",282# "files/museum.jpg",283# "files/terra.jpg",284# "files/underwater.jpg",285# "files/venue.jpg",286# ]287 288 289# with gr.Blocks(title=title, css=custom_css) as demo:290# gr.Markdown(description + article)291 292# # input and control components293# with gr.Row():294# with gr.Column():295# image_input = gr.Image(type='pil', label="Original Image")296# _ = gr.Examples(examples=examples, inputs=[image_input])297# with gr.Column():298# model_dropdown = gr.Dropdown(["vit-small", "vit-large"], label="Model", value="vit-large")299 300# with gr.Accordion('Advanced options (beta)', open=True):301# with gr.Row():302# sensor_width = gr.Number(None, label="Sensor Width in mm", precision=2)303# sensor_height = gr.Number(None, label="Sensor Height in mm", precision=2)304# focal_len = gr.Number(None, label="Focal Length in mm", precision=2)305# camera_detector = gr.Button("Detect Camera from EXIF", elem_id="#button1")306# with gr.Row():307# fx = gr.Number(1000.0, label="fx in pixels", precision=2)308# fy = gr.Number(1000.0, label="fy in pixels", precision=2)309# focal_detector = gr.Button("Calculate Intrinsic", elem_id="#button2")310 311# message_box = gr.Textbox(label="Messages")312 313# # depth and normal314# submit_button = gr.Button("Predict Depth & Normal")315# with gr.Row():316# with gr.Column():317# depth_output = gr.Image(label="Output Depth")318# depth_file = gr.File(label="Depth (.npy)")319# with gr.Column():320# normal_output = gr.Image(label="Output Normal")321# normal_file = gr.File(label="Normal (.npy)")322 323# # 3D reconstruction324# reconstruct_button = gr.Button("Reconstruct 3D")325# pcd_output = gr.Plot(label="3D Point Cloud (Sampled sparse version)")326# pcd_ply = gr.File(label="3D Point Cloud (.ply)")327 328# # cache for depth, normal maps and other states329# state_cache = gr.State({})330 331# # detect focal length in pixels332# camera_detector.click(fn=get_camera, inputs=[image_input], outputs=[sensor_width, sensor_height, focal_len, message_box])333# focal_detector.click(fn=get_intrinsic, inputs=[image_input, sensor_width, sensor_height, focal_len], outputs=[fx, fy, message_box])334 335# submit_button.click(fn=predict_depth_normal, inputs=[image_input, model_dropdown, fx, fy, state_cache], outputs=[depth_output, depth_file, normal_output, normal_file, state_cache, message_box])336# reconstruct_button.click(fn=unprojection_pcd, inputs=[state_cache], outputs=[pcd_output, pcd_ply, message_box])337 338#demo.launch(server_name="0.0.0.0")339 340 341# iface = gr.Interface(342# depth_normal, 343# inputs=[344# gr.Image(type='pil', label="Original Image"),345# gr.Dropdown(["vit-small", "vit-large"], label="Model", info="Select a model type", value="vit-large")346# ],347# outputs=[348# gr.Image(type="pil", label="Output Depth"),349# gr.Image(type="pil", label="Output Normal"),350# gr.Textbox(label="Messages")351# ],352# title=title,353# description=description,354# article=article,355# examples=examples,356# analytics_enabled=False357# )358 359# iface.launch()360 361 362gradio_app = gr.Interface(363 fn=predict_depth_normal,364 inputs=[365 gr.Image(type='pil', label="Original Image"),366 gr.Dropdown(["vit-small", "vit-large"], label="Model"),367 gr.Number(1000.0, label="fx in pixels"),368 gr.Number(1000.0, label="fy in pixels")369 ],370 outputs=[371 gr.Image(label="Output Depth"),372 gr.File(label="Depth (.npy)"),373 gr.Image(label="Output Normal"),374 gr.File(label="Normal (.npy)"),375 gr.Textbox(label="Messages")376 ],377 title="Metric3D",378 description="Metric3Dv2: A versatile monocular geometric foundation model for zero-shot metric depth and surface normal estimation."379)380 381if __name__ == "__main__":382 gradio_app.launch(share=True)