# 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 os import cv2 import torch import gradio as gr from estimator.models.patchfusion import PatchFusion from gradio_imageslider import ImageSlider from torchvision import transforms import numpy as np import torch.nn.functional as F from PIL import Image import tempfile # import spaces import matplotlib from huggingface_hub import hf_hub_download import copy from zoedepth.models.zoedepth import ZoeDepth css = """ #img-display-container { max-height: 100vh; } #img-display-input { max-height: 80vh; } #img-display-output { max-height: 80vh; } """ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' model_name = 'Zhyever/patchfusion_depth_anything_vitl14' # 'Zhyever/patchfusion_depth_anything_vits14', # 'Zhyever/patchfusion_depth_anything_vitb14', # 'Zhyever/patchfusion_depth_anything_vitl14', # 'Zhyever/patchfusion_zoedepth' model = PatchFusion.from_pretrained(model_name).to(DEVICE).eval() pf_ckp = hf_hub_download(repo_id="Zhyever/PatchFusion", filename="DepthAnything_vitl.pt") depth_anything_base = ZoeDepth.build(**model.config.coarse_branch) depth_anything_base.load_state_dict(torch.load(pf_ckp, map_location='cpu')['model']) depth_anything_base = depth_anything_base.to(DEVICE).eval() # @spaces.GPU @torch.no_grad() def predict_depth(model, image_lr, image_hr, mode, patch_number, resolution_h, resolution_w, patch_split_number_h, patch_split_number_w): process_num = 4 tile_cfg = dict() tile_cfg['image_raw_shape'] = (resolution_h, resolution_w) tile_cfg['patch_split_num'] = (patch_split_number_h, patch_split_number_w) if mode == 'r': mode = mode + str(patch_number) result, log_dict = model(mode='infer', cai_mode=mode, process_num=process_num, image_lr=image_lr, image_hr=image_hr, tile_cfg=tile_cfg) return result # def colorize_depth(depth): # depth = (depth - depth.min()) / (depth.max() - depth.min()) # depth = (1 - depth) * 255.0 # depth = depth.astype(np.uint8) # colored_depth = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)[:, :, ::-1] # return colored_depth def colorize_depth(depth_map, min_depth=0, max_depth=0, cmap='Spectral'): """ Colorize depth maps. """ percentile = 0.03 # percentile = 5 min_depth = np.percentile(depth_map, percentile) max_depth = np.percentile(depth_map, 100 - percentile) assert len(depth_map.shape) >= 2, "Invalid dimension" if isinstance(depth_map, torch.Tensor): depth = depth_map.detach().clone().squeeze().numpy() elif isinstance(depth_map, np.ndarray): depth = depth_map.copy().squeeze() # reshape to [ (B,) H, W ] if depth.ndim < 3: depth = depth[np.newaxis, :, :] # colorize cm = matplotlib.colormaps[cmap] depth = ((depth - min_depth) / (max_depth - min_depth)).clip(0, 1) img_colored_np = cm(depth, bytes=False)[:,:,:,0:3] # value from 0 to 1 # img_colored_np = np.rollaxis(img_colored_np, 3, 1) img_colored_np = img_colored_np * 255.0 img_colored_np = img_colored_np.astype(np.uint8) img_colored_np = img_colored_np[0, :, :, :] return img_colored_np def gallery_fn(image, coarse, fine, **kwargs): return [(coarse, fine), image] title = "# PatchFusion" description = """Official demo for **PatchFusion: An End-to-End Tile-Based Framework for High-Resolution Monocular Metric Depth Estimation**. PatchFusion is a deep learning model for high-resolution metric depth estimation from a single image. Please refer to our [project webpage](https://zhyever.github.io/patchfusion), [paper](https://arxiv.org/abs/2312.02284) or [github](https://github.com/zhyever/PatchFusion) for more details. We use Depth-Anything VitL14 as the default model in this demo. You can slide the output to compare the depth prediction from the base Depth-Anything model and PatchFusion. """ markdown = """ PatchFusion works on default 4k (2160x3840) resolution. All input images will be resized to 4k before passing through PatchFusion as default. Users can increase the processing resolution in the advanced option (resolution_h, resolution_w). PatchFusion has three modes: m1, m2, and rn. They are corresponding to the different tiling strategies P16, P49, and Rn in our paper. We rename it for general usage. Please refer to the documentation provided in [github inference instruction](https://github.com/zhyever/PatchFusion/docs/user_infer.md). We provide customized tiling variations for m1 and m2. Users can adjust the patch split number in the advanced option (patch_split_number_h, patch_split_number_w) (default: 4x4 splitting). The output depth map is resized to the original image resolution. Download for better visualization quality. 16-Bit Raw Depth = (pred_depth * 256).to(uint16). """ image_resizer = model.resizer with gr.Blocks(css=css) as demo: gr.Markdown(title) gr.Markdown(description) gr.Markdown("### Depth Prediction demo") gr.Markdown(markdown) prediction_coarse = gr.File( visible=False, ) prediction_fine = gr.File( visible=False, ) with gr.Row(): with gr.Accordion("Advanced options", open=False): resolution_h = gr.Slider(label="Proccessing resolution height (Default 4K, 2160)", minimum=256, maximum=2700, value=2160, step=10) resolution_w = gr.Slider(label="Proccessing resolution width (Default 4K, 3840)", minimum=256, maximum=4800, value=3840, step=10) mode = gr.Radio(["m2", "r"], label="Tiling mode", info="We recommand using M2 for fast evaluation and R with 256 patches for best visualization results, respectively", elem_id='mode', value='m2') # mode = gr.Radio(["m1", "m2", "r"], label="Tiling mode", info="We recommand using M2 for fast evaluation and R with 256 patches for best visualization results, respectively", elem_id='mode', value='m1') patch_number = gr.Slider(1, 256, label="Please decide the number of random patches (Only useful in mode=R)", step=1, value=256) patch_split_number_h = gr.Slider(label="Patch split number of height dimension (Default 4)", minimum=2, maximum=8, value=4, step=1) patch_split_number_w = gr.Slider(label="Patch split number of width dimension (Default 4)", minimum=2, maximum=8, value=4, step=1) color_map = gr.Radio(["turbo_r", "magma_r", "Spectral", "gray"], label="Colormap used to render depth map", elem_id='mode', value='turbo_r') with gr.Row(): input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input') with gr.Row(): depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5) raw_file = gr.File(label="16-Bit Raw Depth, Multiplier:256") submit = gr.Button("Submit") def on_submit(image, mode, patch_number, resolution_h, resolution_w, patch_split_number_h, patch_split_number_w, color_map): h, w = image.shape[:2] image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0 image = transforms.ToTensor()(np.asarray(image)) # raw image image_lr = image_resizer(image.unsqueeze(dim=0)).float().to(DEVICE) # make sure the resolution and patch number are valid # adjust the relationship if resolution_h % (2 * patch_split_number_h) != 0: resolution_h = resolution_h // (2 * patch_split_number_h) * (2 * patch_split_number_h) if resolution_w % (2 * patch_split_number_w) != 0: resolution_w = resolution_w // (2 * patch_split_number_w) * (2 * patch_split_number_w) image_hr = F.interpolate(image.unsqueeze(dim=0), (resolution_h, resolution_w), mode='bicubic', align_corners=True).float().to(DEVICE) pf_prediction = predict_depth(model, image_lr, image_hr, mode, patch_number, resolution_h, resolution_w, patch_split_number_h, patch_split_number_w) coarse_prediction = depth_anything_base(image_lr)['metric_depth'] pf_prediction = F.interpolate(pf_prediction, (h, w))[0, 0].detach().cpu().numpy() coarse_prediction = F.interpolate(coarse_prediction, (h, w))[0, 0].detach().cpu().numpy() pf_prediction_colored = colorize_depth(pf_prediction, cmap=color_map) coarse_predictioncolored = colorize_depth(coarse_prediction, cmap=color_map) raw_depth = Image.fromarray((pf_prediction*256).astype('uint16')) tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False) raw_depth.save(tmp.name) return [(coarse_predictioncolored, pf_prediction_colored), tmp.name] ips = [input_image, mode, patch_number, resolution_h, resolution_w, patch_split_number_h, patch_split_number_w, color_map] submit.click(on_submit, inputs=ips, outputs=[depth_image_slider, raw_file]) inputs = [input_image, prediction_coarse, prediction_fine] outputs = [depth_image_slider, input_image,] examples = gr.Examples( inputs=inputs, outputs=outputs, examples=[ [ "examples/example_1.jpeg", "examples/example_1_coarse.png", "examples/example_1_fine.png", ], [ "examples/example_2.jpeg", "examples/example_2_coarse.png", "examples/example_2_fine.png", ], [ "examples/example_3.jpeg", "examples/example_3_coarse.png", "examples/example_3_fine.png", ]], cache_examples=True, fn=gallery_fn,) if __name__ == '__main__': demo.queue().launch(share=True)