Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import matplotlib
 import numpy as np
 import os
+import time
 from PIL import Image
 import torch
 import torch.nn.functional as F
@@ -11,6 +12,7 @@ import trimesh
 import tempfile
 import shutil
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
 from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
 
@@ -93,6 +95,18 @@ def main(share=True):
         metric_depth[~mask] = metric_depth[mask].max()
         return metric_depth, mask, intrinsics
 
+    def delete_later(path: Union[str, os.PathLike], delay: int = 300):
+        def _delete():
+            try: 
+                os.remove(path) 
+            except FileNotFoundError:
+                pass
+        def _wait_and_delete():
+            time.sleep(delay)
+            _delete(path)
+        thread_pool_executor.submit(_wait_and_delete)
+        atexit.register(_delete)
+
     def on_submit(image, denoise_steps, apply_filter, request: gr.Request = None):
 
         H, W = image.shape[:2]
@@ -157,21 +171,19 @@ def main(share=True):
         vis_path = os.path.join(output_path, '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]
         file_names = ["image_depth_vis.png", "raw_depth.npy", "pointcloud.ply"]
+
+        for f in file_names:
+            delete_later(output_path / f)
+            
         download_files = [
             (output_path / name).as_posix()
             for name in file_names
             if (output_path / name).exists()
         ]
-        # download_files = [vis_path, raw_depth_path, pointcloud_path]
 
         return [(image, colored_depth), glb_path, download_files]
 
-        # return [(image, colored_depth), tmp_glb.name, tmp_concat.name, tmp_raw_depth.name, tmp_ply.name]
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown(title)