add app file for gradio

app.py

@@ -0,0 +1,111 @@
+import os
+import random
+
+import numpy as np
+import torch
+import gradio as gr
+import matplotlib as mpl
+import matplotlib.cm as cm
+
+from vidar.core.wrapper import Wrapper
+from vidar.utils.config import read_config
+
+
+def colormap_depth(depth_map):
+    # Input: depth_map -> HxW numpy array with depth values 
+    # Output: colormapped_im -> HxW numpy array with colorcoded depth values
+    mask = depth_map!=0
+    disp_map = 1/depth_map
+    vmax = np.percentile(disp_map[mask], 95)
+    vmin = np.percentile(disp_map[mask], 5)
+    normalizer = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
+    mapper = cm.ScalarMappable(norm=normalizer, cmap='magma')
+    mask = np.repeat(np.expand_dims(mask,-1), 3, -1)
+    colormapped_im = (mapper.to_rgba(disp_map)[:, :, :3] * 255).astype(np.uint8)
+    colormapped_im[~mask] = 255
+    return colormapped_im
+
+def data_to_batch(data):
+    batch = data.copy()
+    batch['rgb'][0] = batch['rgb'][0].unsqueeze(0).unsqueeze(0)
+    batch['rgb'][1] = batch['rgb'][1].unsqueeze(0).unsqueeze(0)
+    batch['intrinsics'][0] = batch['intrinsics'][0].unsqueeze(0).unsqueeze(0)
+    batch['pose'][0] = batch['pose'][0].unsqueeze(0).unsqueeze(0)
+    batch['pose'][1] = batch['pose'][1].unsqueeze(0).unsqueeze(0)
+    batch['depth'][0] = batch['depth'][0].unsqueeze(0).unsqueeze(0)
+    batch['depth'][1] = batch['depth'][1].unsqueeze(0).unsqueeze(0)
+
+    return batch
+
+
+os.environ['DIST_MODE'] = 'gpu' if torch.cuda.is_available() else 'cpu'
+cfg_file_path = 'configs/papers/define/scannet_temporal_test_context_1.yaml'
+cfg = read_config(cfg_file_path)
+
+wrapper = Wrapper(cfg, verbose=True)
+
+# print('arch: ', wrapper.arch)
+# print('datasets: ', wrapper.datasets)
+
+arch = wrapper.arch
+arch.eval()
+val_dataset = wrapper.datasets['validation'][0]
+len_val_dataset = len(val_dataset)
+# print('val datasets length: ', len_val_dataset)
+
+# data_sample = val_dataset[0]
+# batch = data_to_batch(data_sample)
+# output = arch(batch, epoch=0)
+# print('output: ', output)
+
+# output_depth = output['predictions']['depth'][0][0]
+# print('output_depth: ', output_depth)
+# output_depth = output_depth.squeeze(0).squeeze(0).permute(1,2,0)
+# print('output_depth shape: ', output_depth.shape)
+
+def sample_data_idx():
+    return random.randint(0, len_val_dataset-1)
+
+def display_images_from_idx(idx):
+    rgbs = val_dataset[int(idx)]['rgb']
+    return [np.array(rgb.permute(1,2,0)) for rgb in rgbs.values()]
+
+def infer_depth_from_idx(idx):
+    data_sample = val_dataset[int(idx)]
+    batch = data_to_batch(data_sample)
+    output = arch(batch, epoch=0)
+    output_depths = output['predictions']['depth']
+    return [colormap_depth(output_depth[0].squeeze(0).squeeze(0).squeeze(0).detach().numpy()) for output_depth in output_depths.values()]
+    
+with gr.Blocks() as demo:
+
+    # layout
+    img_box = gr.Gallery(label="Sampled Images").style(grid=[2], height="auto")
+    data_idx_box = gr.Textbox(
+        label="Sampled Data Index", 
+        placeholder="Number between {} and {}".format(0, len_val_dataset-1), 
+        interactive=True
+    )
+    sample_btn = gr.Button('Sample Dataset')
+
+    depth_box = gr.Gallery(label="Infered Depth").style(grid=[2], height="auto")
+    infer_btn = gr.Button('Depth Infer')
+    
+    # actions
+    sample_btn.click(
+        fn=sample_data_idx, 
+        inputs=None, 
+        outputs=data_idx_box
+    ).success(
+        fn=display_images_from_idx,
+        inputs=data_idx_box,
+        outputs=img_box,
+    )
+
+    infer_btn.click(
+        fn=infer_depth_from_idx, 
+        inputs=data_idx_box, 
+        outputs=depth_box
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860)