Update app.py

app.py

@@ -8,6 +8,8 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 from torchvision.transforms import Compose
+import trimesh
+from geometry import create_triangles
 import tempfile
 import spaces
 from zipfile import ZipFile
@@ -159,7 +161,73 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
         # out.release()
         cv2.destroyAllWindows()
         
-        return final_vid, final_zip #output_path
+        return final_vid, final_zip, orig_frames[0], depth_frames[0] #output_path
+
+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 pano_depth_to_world_points(depth):
+    """
+    360 depth to world points
+    given 2D depth is an equirectangular projection of a spherical image
+    Treat depth as radius
+    longitude : -pi to pi
+    latitude : -pi/2 to pi/2
+    """
+
+    # Convert depth to radius
+    radius = depth.flatten()
+
+    lon = np.linspace(-np.pi, np.pi, depth.shape[1])
+    lat = np.linspace(-np.pi/2, np.pi/2, depth.shape[0])
+
+    lon, lat = np.meshgrid(lon, lat)
+    lon = lon.flatten()
+    lat = lat.flatten()
+
+    # Convert to cartesian coordinates
+    x = radius * np.cos(lat) * np.cos(lon)
+    y = radius * np.cos(lat) * np.sin(lon)
+    z = radius * np.sin(lat)
+
+    pts3d = np.stack([x, y, z], axis=1)
+
+    return pts3d
+
+def get_mesh(image, depth, keep_edges=False):
+    image.thumbnail((1024,1024))  # limit the size of the image
+    pts3d = pano_depth_to_world_points(depth)
+
+    # 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 loadurl(url):
     return url
@@ -212,10 +280,16 @@ with gr.Blocks(css=css) as demo:
             input_video = gr.Video(label="Input Video", format="mp4")
             input_url.change(fn=loadurl, inputs=[input_url], outputs=[input_video])
             submit = gr.Button("Submit")
+            render = gr.Button("Render")
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
+            checkbox = gr.Checkbox(label="Keep occlusion edges", value=True)
             processed_video = gr.Video(label="Output Video", format="mp4")
             processed_zip = gr.File(label="Output Archive")
+            output_frame = gr.Image(label="Frame", type='pil')
+            output_depth = gr.Image(label="Depth", type='pil')
+            result = gr.Model3D(label="3D Mesh", clear_color=[
+                                                 1.0, 1.0, 1.0, 1.0])
     
     def on_submit(uploaded_video,model_type):
                 
@@ -224,12 +298,13 @@ with gr.Blocks(css=css) as demo:
 
         return output_video_path
 
-    submit.click(on_submit, inputs=[input_video, model_type], outputs=[processed_video, processed_zip])
+    submit.click(on_submit, inputs=[input_video, model_type], outputs=[processed_video, processed_zip, output_frame, output_depth])
+    render.click(get_mesh, inputs=[output_frame, output_depth, checkbox], outputs=[result])
 
     example_files = os.listdir('examples')
     example_files.sort()
     example_files = [os.path.join('examples', filename) for filename in example_files]
-    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=[processed_video, processed_zip], fn=on_submit, cache_examples=True)
+    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=[processed_video, processed_zip, output_frame, output_depth], fn=on_submit, cache_examples=True)
     
 
 if __name__ == '__main__':