update reconstruction (#7)

- update reconstruct point cloud (f4f7d70b3142cef2c7da97753ad539b2ccc30fd8)

app.py

@@ -29,18 +29,14 @@ from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs,
 import cv2
 from tqdm import tqdm
 import numpy as np
-from PIL import Image
+from PIL import Image, ExifTags
 import matplotlib.pyplot as plt
 
-from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud
+from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud, ply_to_obj
 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
+import plotly.graph_objects as go
 
 #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')
@@ -59,30 +55,27 @@ device = "cuda"
 model_large.to(device)
 model_small.to(device)
 
-
-outputs_dir = "./outs"
-
-def depth_normal(img_path, model_selection="vit-small"):
+def predict_depth_normal(img, model_selection="vit-small", fx=1000.0, fy=1000.0, state_cache={}):
     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
+        return None, None, None, None, state_cache, "Not implemented model."
     
-    img = Image.open(img_path)
+    if img is None:
+        return None, None, None, None, state_cache, "Please upload an image and wait for the upload to complete."
+
 
     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]
+    intrinsic = [fx, fy, 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(
+        pred_depth, pred_depth_scale, scale, output, confidence = get_prediction(
                     model = model,
                     input = rgb_input,
                     cam_model = cam_models_stacks,
@@ -100,125 +93,267 @@ def depth_normal(img_path, model_selection="vit-small"):
     pred_depth = pred_depth.squeeze().cpu().numpy()
     pred_depth[pred_depth<0] = 0
     pred_color = gray_to_colormap(pred_depth)
+
+    pred_normal = torch.nn.functional.interpolate(pred_normal, [img.shape[0], img.shape[1]], mode='bilinear').squeeze()
+    pred_normal = pred_normal.permute(1,2,0)
+    pred_color_normal = vis_surface_normal(pred_normal)
+    pred_normal = pred_normal.cpu().numpy()
     
+    # Storing depth and normal map in state for potential 3D reconstruction
+    state_cache['depth'] = pred_depth
+    state_cache['normal'] = pred_normal
+    state_cache['img'] = img
+    state_cache['intrinsic'] = intrinsic
+    state_cache['confidence'] = confidence 
+
+    # save depth and normal map to .npy file
+    if 'save_dir' not in state_cache:
+        cache_id = np.random.randint(0, 100000000000)
+        while osp.exists(f'recon_cache/{cache_id:08d}'):
+            cache_id = np.random.randint(0, 100000000000)
+        state_cache['save_dir'] = f'recon_cache/{cache_id:08d}'
+        os.makedirs(state_cache['save_dir'], exist_ok=True)
+    depth_file = f"{state_cache['save_dir']}/depth.npy"
+    normal_file = f"{state_cache['save_dir']}/normal.npy"
+    np.save(depth_file, pred_depth)
+    np.save(normal_file, pred_depth)
+
     ##formatted = (output * 255 / np.max(output)).astype('uint8')
+    img = Image.fromarray(pred_color)
+    img_normal = Image.fromarray(pred_color_normal)
+    return img, depth_file, img_normal, normal_file, state_cache, "Success!"
+
+def get_camera(img):
+    if img is None:
+        return None, None, None, "Please upload an image and wait for the upload to complete."
+    try:
+        exif = img.getexif()
+        exif.update(exif.get_ifd(ExifTags.IFD.Exif))
+    except:
+        exif = {}
+    sensor_width = exif.get(ExifTags.Base.FocalPlaneYResolution, None)
+    sensor_height = exif.get(ExifTags.Base.FocalPlaneXResolution, None)
+    focal_length = exif.get(ExifTags.Base.FocalLength, None)
+    
+    # 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-from
+    w, h = img.size
+    sensor_width = w / sensor_width * 25.4 if sensor_width is not None else None
+    sensor_height = h / sensor_height * 25.4 if sensor_height is not None else None
+    focal_length = focal_length * 1.0 if focal_length is not None else None
+
+    message = "Success!"
+    if focal_length is None:
+        message = "Focal length not found in EXIF. Please manually input."
+    elif sensor_width is None and sensor_height is None:
+        sensor_width = 16
+        sensor_height = h / w * sensor_width
+        message = f"Sensor size not found in EXIF. Using {sensor_width}x{sensor_height:.2f} mm as default."
+
+    return sensor_width, sensor_height, focal_length, message
+
+def get_intrinsic(img, sensor_width, sensor_height, focal_length):
+    if img is None:
+        return None, None, "Please upload an image and wait for the upload to complete."
+    if sensor_width is None or sensor_height is None or focal_length is None:
+        return 1000, 1000, "Insufficient information. Try detecting camera first or use default 1000 for fx and fy."
+    if sensor_width == 0 or sensor_height == 0 or focal_length == 0:
+        return 1000, 1000, "Insufficient information. Try detecting camera first or use default 1000 for fx and fy."
     
-    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)
+    # calculate focal length in pixels
+    w, h = img.size
+    fx = w / sensor_width * focal_length if sensor_width is not None else None
+    fy = h / sensor_height * focal_length if sensor_height is not None else None
 
-    name_base = os.path.splitext(os.path.basename(img_path))[0]
+    # if fx is None:
+    #     return fy, fy, "Sensor width not provided, using fy for both fx and fy"
+    # if fy is None:
+    #     return fx, fx, "Sensor height not provided, using fx for both fx and fy"
 
-    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)
+    return fx, fy, "Success!"
 
-    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")
+def unprojection_pcd(state_cache):
+    depth_map = state_cache.get('depth', None)
+    normal_map = state_cache.get('normal', None)
+    img = state_cache.get('img', None)
+    intrinsic = state_cache.get('intrinsic', None)
 
-    np.save(normal_path, normal_np)
-    np.save(depth_path, depth_np)
+    if depth_map is None or img is None:
+        return None, "Please predict depth and normal first."
     
-    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
+    # # downsample/upsample the depth map to confidence map size
+    # confidence = state_cache.get('confidence', None)
+    # if confidence is not None:
+    #     H, W = confidence.shape
+    #     # intrinsic[0] *= W / depth_map.shape[1]
+    #     # intrinsic[1] *= H / depth_map.shape[0]
+    #     # intrinsic[2] *= W / depth_map.shape[1]
+    #     # intrinsic[3] *= H / depth_map.shape[0]
+    #     depth_map = cv2.resize(depth_map, (W, H), interpolation=cv2.INTER_LINEAR)
+    #     img = cv2.resize(img, (W, H), interpolation=cv2.INTER_LINEAR)
+    
+    #     # filter out depth map by confidence
+    #     mask = confidence.cpu().numpy() > 0
+
+    # downsample the depth map if too large
+    if depth_map.shape[0] > 1080:
+        scale = 1080 / depth_map.shape[0]
+        depth_map = cv2.resize(depth_map, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
+        img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
+        intrinsic = [intrinsic[0]*scale, intrinsic[1]*scale, intrinsic[2]*scale, intrinsic[3]*scale]
+    
+    if 'save_dir' not in state_cache:
+        cache_id = np.random.randint(0, 100000000000)
+        while osp.exists(f'recon_cache/{cache_id:08d}'):
+            cache_id = np.random.randint(0, 100000000000)
+        state_cache['save_dir'] = f'recon_cache/{cache_id:08d}'
+        os.makedirs(state_cache['save_dir'], exist_ok=True)
+
+    pcd_ply = f"{state_cache['save_dir']}/output.ply"
+    pcd_obj = pcd_ply.replace(".ply", ".obj")
+    
+    pcd = reconstruct_pcd(depth_map, intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3])
+    # if mask is not None:
+    #     pcd_filtered = pcd[mask]
+    #     img_filtered = img[mask]
+    pcd_filtered = pcd.reshape(-1, 3)
+    img_filtered = img.reshape(-1, 3)
+
+    save_point_cloud(pcd_filtered, img_filtered, pcd_ply, binary=False)
+    # ply_to_obj(pcd_ply, pcd_obj)
+
+    # downsample the point cloud for visualization
+    num_samples = 250000
+    if pcd_filtered.shape[0] > num_samples:
+        indices = np.random.choice(pcd_filtered.shape[0], num_samples, replace=False)
+        pcd_downsampled = pcd_filtered[indices]
+        img_downsampled = img_filtered[indices]
+    else:
+        pcd_downsampled = pcd_filtered
+        img_downsampled = img_filtered
+
+    # plotly show
+    color_str = np.array([f"rgb({r},{g},{b})" for b,g,r in img_downsampled])
+    data=[go.Scatter3d(
+        x=pcd_downsampled[:,0],
+        y=pcd_downsampled[:,1],
+        z=pcd_downsampled[:,2],
+        mode='markers',
+        marker=dict(
+            size=1,
+            color=color_str,
+            opacity=0.8,
+        )
+    )]
+    layout = go.Layout(
+        margin=dict(l=0, r=0, b=0, t=0),
+        scene=dict(
+            camera = dict(
+                eye=dict(x=0, y=0, z=-1),
+                up=dict(x=0, y=-1, z=0)
+            ),
+            xaxis=dict(showgrid=False, showticklabels=False, visible=False),
+            yaxis=dict(showgrid=False, showticklabels=False, visible=False),
+            zaxis=dict(showgrid=False, showticklabels=False, visible=False),
+        )
+    )
+    fig = go.Figure(data=data, layout=layout)
+
+    return fig, pcd_ply, "Success!"
 
-    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 = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2307.10984.pdf'>Metric3D: Towards Zero-shot Metric 3D Prediction from A Single Image</a> | <a href='https://github.com/YvanYin/Metric3D'>Github Repo</a></p>"
+description = '''# Metric3Dv2: A versatile monocular geometric foundation model for zero-shot metric depth and surface normal estimation
+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.'''
+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>"
+
+custom_css = '''#button1, #button2 {
+    width: 20px;
+}'''
 
 examples = [
-    ["files/museum.jpg"],
-    ["files/terra.jpg"],
-    ["files/underwater.jpg"],
-    ["files/venue.jpg"]
+    #["turtle.jpg"],
+    #["lions.jpg"]
+    #["files/gundam.jpg"],
+    "files/p50_pro.jpg",
+    "files/iphone13.JPG",
+    "files/canon_cat.JPG",
+    "files/canon_dog.JPG",
+    "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)
+with gr.Blocks(title=title, css=custom_css) as demo:
+    gr.Markdown(description + article)
+
+    # input and control components
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type='pil', label="Original Image")
+            _ = gr.Examples(examples=examples, inputs=[image_input])
+        with gr.Column():
+            model_dropdown = gr.Dropdown(["vit-small", "vit-large"], label="Model", value="vit-large")
+
+            with gr.Accordion('Advanced options (beta)', open=True):
+                with gr.Row():
+                    sensor_width = gr.Number(None, label="Sensor Width in mm", precision=2)
+                    sensor_height = gr.Number(None, label="Sensor Height in mm", precision=2)
+                    focal_len = gr.Number(None, label="Focal Length in mm", precision=2)
+                    camera_detector = gr.Button("Detect Camera from EXIF", elem_id="#button1")
+                with gr.Row():
+                    fx = gr.Number(1000.0, label="fx in pixels", precision=2)
+                    fy = gr.Number(1000.0, label="fy in pixels", precision=2)
+                    focal_detector = gr.Button("Calculate Intrinsic", elem_id="#button2")
+
+            message_box = gr.Textbox(label="Messages")
+
+    # depth and normal
+    submit_button = gr.Button("Predict Depth and Normal")
+    with gr.Row():
+        with gr.Column():
+            depth_output = gr.Image(label="Output Depth")
+            depth_file = gr.File(label="Depth (.npy)")
+        with gr.Column():
+            normal_output = gr.Image(label="Output Normal")
+            normal_file = gr.File(label="Normal (.npy)")
+
+    # 3D reconstruction
+    reconstruct_button = gr.Button("Reconstruct 3D")
+    pcd_output = gr.Plot(label="3D Point Cloud (Sampled sparse version)")
+    pcd_ply = gr.File(label="3D Point Cloud (.ply)")
+
+    # cache for depth, normal maps and other states
+    state_cache = gr.State({})
+
+    # detect focal length in pixels
+    camera_detector.click(fn=get_camera, inputs=[image_input], outputs=[sensor_width, sensor_height, focal_len, message_box])
+    focal_detector.click(fn=get_intrinsic, inputs=[image_input, sensor_width, sensor_height, focal_len], outputs=[fx, fy, message_box])
+
+    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])
+    reconstruct_button.click(fn=unprojection_pcd, inputs=[state_cache], outputs=[pcd_output, pcd_ply, message_box])
+
+demo.launch()
+
+
+# iface = gr.Interface(
+#     depth_normal, 
+#     inputs=[
+#         gr.Image(type='pil', label="Original Image"),
+#         gr.Dropdown(["vit-small", "vit-large"], label="Model", info="Select a model type", value="vit-large")
+#     ],
+#     outputs=[
+#         gr.Image(type="pil", label="Output Depth"),
+#         gr.Image(type="pil", label="Output Normal"),
+#         gr.Textbox(label="Messages")
+#     ],
+#     title=title,
+#     description=description,
+#     article=article,
+#     examples=examples,
+#     analytics_enabled=False
+# )
+
+# iface.launch()

mono/utils/do_test.py

@@ -154,9 +154,10 @@ def get_prediction(
     )
     #pred_depth, confidence, output_dict = model.module.inference(data)
     pred_depth, confidence, output_dict = model.inference(data)
-    pred_depth = pred_depth
     pred_depth = pred_depth.squeeze()
     pred_depth = pred_depth[pad_info[0] : pred_depth.shape[0] - pad_info[1], pad_info[2] : pred_depth.shape[1] - pad_info[3]]
+    confidence = confidence.squeeze()
+    confidence = confidence[pad_info[0] : confidence.shape[0] - pad_info[1], pad_info[2] : confidence.shape[1] - pad_info[3]]
     if gt_depth is not None:
         resize_shape = gt_depth.shape
     elif ori_shape != []:
@@ -172,7 +173,7 @@ def get_prediction(
         pred_depth_scale = None
         scale = None
 
-    return pred_depth, pred_depth_scale, scale, output_dict
+    return pred_depth, pred_depth_scale, scale, output_dict, confidence
 
 def transform_test_data_scalecano(rgb, intrinsic, data_basic):
     """

mono/utils/unproj_pcd.py

@@ -3,6 +3,7 @@ import torch
 from plyfile import PlyData, PlyElement
 import cv2
 
+import trimesh
 
 def get_pcd_base(H, W, u0, v0, fx, fy):
     x_row = np.arange(0, W)
@@ -85,4 +86,32 @@ def save_point_cloud(pcd, rgb, filename, binary=True):
                     'property uchar blue\n' \
                     'end_header' % r.shape[0]
         # ---- Save ply data to disk
-        np.savetxt(filename, np.column_stack[x, y, z, r, g, b], fmt='%f %f %f %d %d %d', header=ply_head, comments='')
+        np.savetxt(filename, np.column_stack([x, y, z, r, g, b]), fmt='%f %f %f %d %d %d', header=ply_head, comments='')
+
+def ply_to_obj(ply_file, obj_file):
+    mesh = trimesh.load_mesh(ply_file)
+    # T2 = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
+    # mesh.apply_transform(T2)
+    mesh.export(obj_file)
+
+
+# import numpy as np
+
+# def save_point_cloud_to_obj(points, colors, file_name):
+#     """
+#     Save a numpy array of point cloud data with color to an OBJ file.
+    
+#     Args:
+#     points (np.ndarray): A numpy array of shape (H, W, 3) where H is height, W is width.
+#     colors (np.ndarray): A numpy array of color data, shape (H, W, 3), values should be in [0, 1].
+#     file_name (str): The path to the output .obj file.
+#     """
+#     H, W, _ = points.shape
+#     assert points.shape == colors.shape, "Points and colors must have the same shape"
+    
+#     with open(file_name, 'w') as file:
+#         for i in range(H):
+#             for j in range(W):
+#                 x, y, z = points[i, j]
+#                 r, g, b = colors[i, j]
+#                 file.write(f"v {x} {y} {z} {r} {g} {b}\n")

requirements.txt

@@ -19,4 +19,6 @@ datetime
 gradio==4.11.0
 gradio-imageslider==0.0.16
 cupy-cuda12x
-plotly
+plotly
+trimesh
+exifread