app.py

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, ExifTags
import matplotlib.pyplot as plt

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 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')

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)

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:
        return None, None, None, None, state_cache, "Not implemented model."
    
    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 = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
    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, confidence = 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)

    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."
    
    # 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

    # 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"

    return fx, fy, "Success!"


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)

    if depth_map is None or img is None:
        return None, "Please predict depth and normal first."
    
    # # 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!"


title = "Metric3D"
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 = [
    #["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",
]


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()