ui_prediction.py

# MIT License

# Copyright (c) 2022 Intelligent Systems Lab Org

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# File author: Zhenyu Li

import gradio as gr
from PIL import Image
import tempfile
import torch
import numpy as np

from zoedepth.utils.arg_utils import parse_unknown
import argparse
from zoedepth.models.builder import build_model
from zoedepth.utils.config import get_config_user
import matplotlib
import cv2

from infer_user import regular_tile_param, random_tile_param
from zoedepth.models.base_models.midas import Resize
from torchvision.transforms import Compose
from PIL import Image
from torchvision import transforms
import torch.nn.functional as F

from zoedepth.models.base_models.midas import Resize
from torchvision.transforms import Compose

import gradio as gr
import numpy as np
import trimesh
from zoedepth.utils.geometry import depth_to_points, create_triangles
from functools import partial
import tempfile

def depth_edges_mask(depth, occ_filter_thr):
    """Returns a mask of edges in the depth map.
    Args:
    depth: 2D numpy array of shape (H, W) with dtype float32.
    Returns:
    mask: 2D numpy array of shape (H, W) with dtype bool.
    """
    # Compute the x and y gradients of the depth map.
    depth_dx, depth_dy = np.gradient(depth)
    # Compute the gradient magnitude.
    depth_grad = np.sqrt(depth_dx ** 2 + depth_dy ** 2)
    # Compute the edge mask.
    # mask = depth_grad > 0.05 # default in zoedepth
    mask = depth_grad > occ_filter_thr # preserve more edges (?)
    return mask

def load_state_dict(model, state_dict):
    """Load state_dict into model, handling DataParallel and DistributedDataParallel. Also checks for "model" key in state_dict.

    DataParallel prefixes state_dict keys with 'module.' when saving.
    If the model is not a DataParallel model but the state_dict is, then prefixes are removed.
    If the model is a DataParallel model but the state_dict is not, then prefixes are added.
    """
    state_dict = state_dict.get('model', state_dict)
    # if model is a DataParallel model, then state_dict keys are prefixed with 'module.'

    do_prefix = isinstance(
        model, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel))
    state = {}
    for k, v in state_dict.items():
        if k.startswith('module.') and not do_prefix:
            k = k[7:]

        if not k.startswith('module.') and do_prefix:
            k = 'module.' + k

        state[k] = v

    model.load_state_dict(state, strict=True)
    print("Loaded successfully")
    return model

def load_wts(model, checkpoint_path):
    ckpt = torch.load(checkpoint_path, map_location='cpu')
    return load_state_dict(model, ckpt)

def load_ckpt(model, checkpoint):
    model = load_wts(model, checkpoint)
    print("Loaded weights from {0}".format(checkpoint))
    return model

def colorize(value, cmap='magma_r', vmin=None, vmax=None):
    # normalize
    vmin = value.min() if vmin is None else vmin
    # vmax = value.max() if vmax is None else vmax
    vmax = np.percentile(value, 95) if vmax is None else vmax

    if vmin != vmax:
        value = (value - vmin) / (vmax - vmin)  # vmin..vmax
    else:
        value = value * 0.

    cmapper = matplotlib.cm.get_cmap(cmap)
    value = cmapper(value, bytes=True)  # ((1)xhxwx4)

    value = value[:, :, :3] # bgr -> rgb
    # rgb_value = value[..., ::-1]
    rgb_value = value

    return rgb_value

def predict_depth(model, image, mode, pn, reso, ps, device=None):

    pil_image = image
    if device is not None:
        image = transforms.ToTensor()(pil_image).unsqueeze(0).to(device)
    else:
        image = transforms.ToTensor()(pil_image).unsqueeze(0).cuda()
        
    image_height, image_width = image.shape[-2], image.shape[-1]

    if reso != '':
        image_resolution = (int(reso.split('x')[0]), int(reso.split('x')[1]))
    else:
        image_resolution = (2160, 3840)
    image_hr = F.interpolate(image, image_resolution, mode='bicubic', align_corners=True)
    preprocess = Compose([Resize(512, 384, keep_aspect_ratio=False, ensure_multiple_of=32, resize_method="minimal")])
    image_lr = preprocess(image)

    if ps != '':
        patch_size = (int(ps.split('x')[0]), int(ps.split('x')[1]))
    else:
        patch_size = (int(image_resolution[0] // 4), int(image_resolution[1] // 4))

    avg_depth_map = regular_tile_param(
        model, 
        image_hr, 
        offset_x=0, 
        offset_y=0, 
        img_lr=image_lr,
        crop_size=patch_size, 
        img_resolution=image_resolution, 
        transform=preprocess,
        blr_mask=True)

    if mode== 'P16':
        pass
    elif mode== 'P49':
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=patch_size[1]//2, 
            offset_y=0, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=0, 
            offset_y=patch_size[0]//2, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=patch_size[1]//2, 
            offset_y=patch_size[0]//2, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)
    elif mode == 'R':
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=patch_size[1]//2, 
            offset_y=0, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=0, 
            offset_y=patch_size[0]//2, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)
        regular_tile_param(
            model, 
            image_hr, 
            offset_x=patch_size[1]//2, 
            offset_y=patch_size[0]//2, 
            img_lr=image_lr, 
            iter_pred=avg_depth_map.average_map, 
            boundary=0, 
            update=True, 
            avg_depth_map=avg_depth_map, 
            crop_size=patch_size, 
            img_resolution=image_resolution, 
            transform=preprocess,
            blr_mask=True)

        for i in range(int(pn)):
            random_tile_param(
                model, 
                image_hr, 
                img_lr=image_lr, 
                iter_pred=avg_depth_map.average_map, 
                boundary=0, 
                update=True, 
                avg_depth_map=avg_depth_map, 
                crop_size=patch_size, 
                img_resolution=image_resolution, 
                transform=preprocess,
                blr_mask=True)
    
    depth = avg_depth_map.average_map.detach().cpu()
    depth = F.interpolate(depth.unsqueeze(dim=0).unsqueeze(dim=0), (image_height, image_width), mode='bicubic', align_corners=True).squeeze().numpy()

    return depth

def create_demo(model):
    gr.Markdown("## Depth Prediction Demo")

    with gr.Accordion("Advanced options", open=False):
        mode = gr.Radio(["P49", "R"], label="Tiling mode", info="We recommand using P49 for fast evaluation and R with 1024 patches for best visualization results, respectively", elem_id='mode', value='R'),
        patch_number = gr.Slider(1, 1024, label="Please decide the number of random patches (Only useful in mode=R)", step=1, value=256)
        resolution = gr.Textbox(label="Proccessing resolution (Default 4K. Use 'x' to split height and width.)", elem_id='mode', value='2160x3840')
        patch_size = gr.Textbox(label="Patch size (Default 1/4 of image resolution. Use 'x' to split height and width.)", elem_id='mode', value='540x960')
    
    with gr.Row():
        input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')
        depth_image = gr.Image(label="Depth Map", elem_id='img-display-output')
    raw_file = gr.File(label="16-bit raw depth, multiplier:256")
    submit = gr.Button("Submit")

    def on_submit(image, mode, pn, reso, ps):
        depth = predict_depth(model, image, mode, pn, reso, ps)
        colored_depth = colorize(depth, cmap='gray_r')
        tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
        raw_depth = Image.fromarray((depth*256).astype('uint16'))
        raw_depth.save(tmp.name)
        return [colored_depth, tmp.name]
    
    submit.click(on_submit, inputs=[input_image, mode[0], patch_number, resolution, patch_size], outputs=[depth_image, raw_file])
    examples = gr.Examples(examples=["examples/example_1.jpeg", "examples/example_2.jpeg", "examples/example_3.jpeg"], inputs=[input_image])

def get_mesh(model, image, mode, pn, reso, ps, keep_edges, occ_filter_thr, fov):
    depth = predict_depth(model, image, mode, pn, reso, ps)

    image.thumbnail((1024,1024))  # limit the size of the input image
    depth = F.interpolate(torch.from_numpy(depth).unsqueeze(dim=0).unsqueeze(dim=0), (image.height, image.width), mode='bicubic', align_corners=True).squeeze().numpy()

    pts3d = depth_to_points(depth[None], fov=float(fov))
    pts3d = pts3d.reshape(-1, 3)

    # Create a trimesh mesh from the points
    # Each pixel is connected to its 4 neighbors
    # colors are the RGB values of the image

    verts = pts3d.reshape(-1, 3)
    image = np.array(image)
    if keep_edges:
        triangles = create_triangles(image.shape[0], image.shape[1])
    else:
        triangles = create_triangles(image.shape[0], image.shape[1], mask=~depth_edges_mask(depth, occ_filter_thr=float(occ_filter_thr)))
    colors = image.reshape(-1, 3)
    mesh = trimesh.Trimesh(vertices=verts, faces=triangles, vertex_colors=colors)

    # Save as glb
    glb_file = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
    glb_path = glb_file.name
    mesh.export(glb_path)
    return glb_path

def create_demo_3d(model):

    gr.Markdown("### Image to 3D Mesh")
    gr.Markdown("Convert a single 2D image to a 3D mesh")

    with gr.Accordion("Advanced options", open=False):
        mode = gr.Radio(["P49", "R"], label="Tiling mode", info="We recommand using P49 for fast evaluation and R with 1024 patches for best visualization results, respectively", elem_id='mode', value='R'),
        patch_number = gr.Slider(1, 1024, label="Please decide the number of random patches (Only useful in mode=R)", step=1, value=256)
        resolution = gr.Textbox(label="Proccessing resolution (Default 4K. Use 'x' to split height and width)", value='2160x3840')
        patch_size = gr.Textbox(label="Patch size (Default 1/4 of image resolution. Use 'x' to split height and width)", value='540x960')

        checkbox = gr.Checkbox(label="Keep occlusion edges", value=False)
        # occ_filter_thr = gr.Textbox(label="Occlusion filter threshold", info="Larger value will reserve more edges (Only useful when NOT keeping occlusion edges)", value='0.5')
        # fov = gr.Textbox(label="FOV for inv-projection", value='55')

        occ_filter_thr = gr.Slider(0.01, 5, label="Occlusion edge filter threshold", info="Larger value will reserve more occlusion edges (Only useful when NOT keeping occlusion edges)", step=0.01, value=0.2)
        fov = gr.Slider(5, 180, label="FOV for inv-projection", step=1, value=55)


    with gr.Row():
        input_image = gr.Image(label="Input Image", type='pil')
        result = gr.Model3D(label="3d mesh reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
    
    submit = gr.Button("Submit")
    submit.click(partial(get_mesh, model), inputs=[input_image, mode[0], patch_number, resolution, patch_size, checkbox, occ_filter_thr, fov], outputs=[result])
    examples = gr.Examples(examples=["examples/example_1.jpeg", "examples/example_4.jpeg", "examples/example_3.jpeg"], inputs=[input_image])