Spaces:
Runtime error
Runtime error
| # 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: Shariq Farooq Bhat | |
| 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): | |
| """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 | |
| return mask | |
| def predict_depth(model, image): | |
| depth = model.infer_pil(image) | |
| return depth | |
| def get_mesh(model, image, keep_edges=False): | |
| image.thumbnail((1024,1024)) # limit the size of the input image | |
| depth = predict_depth(model, image) | |
| pts3d = depth_to_points(depth[None]) | |
| 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)) | |
| 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(model): | |
| gr.Markdown("### Image to 3D mesh") | |
| gr.Markdown("Convert a single 2D image to a 3D mesh") | |
| with gr.Row(): | |
| 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]) | |
| checkbox = gr.Checkbox(label="Keep occlusion edges", value=False) | |
| submit = gr.Button("Submit") | |
| submit.click(partial(get_mesh, model), inputs=[image, checkbox], outputs=[result]) | |
| # examples = gr.Examples(examples=["examples/aerial_beach.jpeg", "examples/mountains.jpeg", "examples/person_1.jpeg", "examples/ancient-carved.jpeg"], | |
| # inputs=[image]) | |