update

ui_prediction.py

@@ -0,0 +1,347 @@
+# 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])