Update app.py

app.py

@@ -112,26 +112,34 @@ def main(share=True):
         # metric depth -> point cloud
         pcd = depth2pcd(metric_depth, intrinsics, color=cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB), input_mask=mask, ret_pcd=True)
         if apply_filter:
-            cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=3.0)
+            cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
             pcd = pcd.select_by_index(ind)
 
+        tempdir = tempfile.TemporaryDirectory()
+        out_dir = tempdir.name
+        pointcloud_path = os.path.join(out_dir, 'pointcloud.ply')
+
         # save pcd to temporary .ply
         tmp_ply = tempfile.NamedTemporaryFile(suffix='.ply', delete=False)
-        o3d.io.write_point_cloud(tmp_ply.name, pcd)
+        o3d.io.write_point_cloud(pointcloud_path, pcd)
         
         vertices = np.asarray(pcd.points)
         vertex_colors = (np.asarray(pcd.colors) * 255).astype(np.uint8)
         vertices[:, 2] = -vertices[:, 2]
         vertices[:, 1] = -vertices[:, 1]
         mesh = trimesh.PointCloud(vertices=vertices, colors=vertex_colors)
-        tmp_glb = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
+        glb_path = os.path.join(out_dir, 'pointcloud.glb')
+        # tmp_glb = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
         mesh.export(tmp_glb.name)
 
 
         # save raw depth (npy)
         depth = cv2.resize(ppd_depth, (W, H), interpolation=cv2.INTER_LINEAR)
-        tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.npy', delete=False)
-        np.save(tmp_raw_depth.name, depth)
+        raw_depth_path = os.path.join(out_dir, 'raw_depth.npy')
+        np.save(raw_depth_path, depth)
+    
+        # tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.npy', delete=False)
+        # np.save(tmp_raw_depth.name, depth)
 
         depth_vis = (depth - depth.min()) / (depth.max() - depth.min() + 1e-5) * 255.0
         depth_vis = depth_vis.astype(np.uint8)
@@ -139,10 +147,19 @@ def main(share=True):
 
         split_region = np.ones((image.shape[0], 50, 3), dtype=np.uint8) * 255
         combined_result = cv2.hconcat([image[:, :, ::-1], split_region, colored_depth[:, :, ::-1]])
-        tmp_concat = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-        cv2.imwrite(tmp_concat.name, combined_result)
 
-        return [(image, colored_depth), tmp_glb.name, tmp_concat.name, tmp_raw_depth.name, tmp_ply.name]
+        vis_path = os.path.join(out_dir, 'image_depth_vis.png')
+        cv2.imwrite(vis_path, combined_result)
+    
+        # tmp_concat = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+        # cv2.imwrite(tmp_concat.name, combined_result)
+
+        # download_files = [tmp_concat.name, tmp_raw_depth.name, tmp_ply.name]
+        download_files = [vis_path, raw_depth_path, pointcloud_path]
+
+        return [(image, colored_depth), tmp_glb.name, download_files]
+
+        # return [(image, colored_depth), tmp_glb.name, tmp_concat.name, tmp_raw_depth.name, tmp_ply.name]
 
     with gr.Blocks(css=css) as demo:
         gr.Markdown(title)
@@ -156,25 +173,29 @@ def main(share=True):
                 with gr.Accordion(label="Settings", open=False):
                     denoise_steps = gr.Slider(label="Denoising Steps", minimum=1, maximum=50, value=10, step=1)
                     apply_filter = gr.Checkbox(label="Apply filter points", value=True)
-                submit_btn = gr.Button(value="Predict Depth")
+                submit_btn = gr.Button(value="Predict")
 
             # Right: 3D point cloud + depth
             with gr.Column():
                 with gr.Tabs():
-                    with gr.Tab("3D View"):
-                        model_3d = gr.Model3D(display_mode="solid", label="3D Point Map", clear_color=[1,1,1,1], height="60vh")
                     with gr.Tab("Depth"):
                         depth_map = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+                    with gr.Tab("3D View"):
+                        model_3d = gr.Model3D(display_mode="solid", label="3D Point Map", clear_color=[1,1,1,1], height="60vh")
+                    with gr.Tab("Download"):
+                        download_files = gr.File(type='file', label="Download Files", file_types=[".ply", ".glb", ".npy"], file_types_mode="multiple")
     
-        concat_file = gr.File(label="Concatenated visualization (image+depth)", elem_id="image-depth-download")
-        raw_depth_file = gr.File(label="Raw depth output (saved as .npy)", elem_id="download")
-        pcd_file = gr.File(label="Point Cloud (.ply)", elem_id="download-ply")
+        # concat_file = gr.File(label="Concatenated visualization (image+depth)", elem_id="image-depth-download")
+        # raw_depth_file = gr.File(label="Raw depth output (saved as .npy)", elem_id="download")
+        # pcd_file = gr.File(label="Point Cloud (.ply)", elem_id="download-ply")
 
         submit_btn.click(
-            on_submit, 
-            inputs=[input_image, denoise_steps, apply_filter], 
-            outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file]
-        )
+            fn=lambda: [None, None, None, "", "", ""],
+            outputs=[depth_map, model_3d, download_files]
+        ).then(
+            fn=on_submit,
+            inputs=[input_image, denoise_steps, apply_filter],
+            outputs=[depth_map, model_3d, download_files]
 
         example_files = os.listdir('assets/examples')
         example_files.sort()
@@ -182,7 +203,7 @@ def main(share=True):
         examples = gr.Examples(
             examples=example_files, 
             inputs=[input_image], 
-            outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file], 
+            outputs=[depth_map, model_3d, download_files], 
             fn=on_submit
         )