Spaces:
Running
Running
| import gradio as gr | |
| import spaces | |
| import tempfile | |
| import os | |
| import torch | |
| import numpy as np | |
| from matplotlib.colors import LinearSegmentedColormap | |
| from app_config import CSS, TITLE, DESCRIPTION, DEVICE | |
| import sample_cond | |
| model = sample_cond.load_model() | |
| def create_custom_colormap(): | |
| colors = [(0, 1, 0), (0, 1, 1), (0, 0, 1), (1, 0, 1), (1, 1, 0)] | |
| positions = [0, 0.38, 0.6, 0.7, 1] | |
| custom_cmap = LinearSegmentedColormap.from_list('custom_colormap', list(zip(positions, colors)), N=256) | |
| return custom_cmap | |
| def colorize_depth(depth, log_scale): | |
| if log_scale: | |
| depth = ((np.log2((depth / 255.) * 56. + 1) / 5.84) * 255.).astype(np.uint8) | |
| mask = depth == 0 | |
| colormap = create_custom_colormap() | |
| rgb = colormap(depth)[:, :, :3] | |
| rgb[mask] = 0. | |
| return rgb | |
| def generate_lidar(model, cond): | |
| img, pcd = sample_cond.sample(model, cond) | |
| return img, pcd | |
| def load_camera(image): | |
| split_per_view = 4 | |
| camera = np.array(image).astype(np.float32) / 255. | |
| camera = camera.transpose(2, 0, 1) | |
| camera_list = np.split(camera, split_per_view, axis=2) # split into n chunks as different views | |
| camera_cond = torch.from_numpy(np.stack(camera_list, axis=0)).unsqueeze(0).to(DEVICE) | |
| return camera_cond | |
| with gr.Blocks(css=CSS) as demo: | |
| gr.Markdown(TITLE) | |
| gr.Markdown(DESCRIPTION) | |
| gr.Markdown("### Camera-to-LiDAR Demo") | |
| # gr.Markdown("You can slide the output to compare the depth prediction with input image") | |
| with gr.Row(): | |
| input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input') | |
| output_image = gr.Image(label="Range Map", elem_id='img-display-output') | |
| raw_file = gr.File(label="Point Cloud (.txt file). Can be viewed through Meshlab") | |
| submit = gr.Button("Submit") | |
| def on_submit(image): | |
| cond = load_camera(image) | |
| img, pcd = generate_lidar(model, cond) | |
| tmp = tempfile.NamedTemporaryFile(suffix='.txt', delete=False) | |
| pcd.save(tmp.name) | |
| rgb_img = colorize_depth(img, log_scale=True) | |
| return [rgb_img, tmp.name] | |
| submit.click(on_submit, inputs=[input_image], outputs=[output_image, raw_file]) | |
| example_files = sorted(os.listdir('cam_examples')) | |
| example_files = [os.path.join('cam_examples', filename) for filename in example_files] | |
| examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[output_image, raw_file], | |
| fn=on_submit, cache_examples=True) | |
| if __name__ == '__main__': | |
| demo.queue().launch() | |