inital version

app.py

@@ -1,13 +1,11 @@
 import os
-# import time
-# import shutil
+import time
+import shutil
 from pathlib import Path
-# from typing import Union
-# import atexit
+from typing import Union
+import atexit
 import spaces
-# from concurrent.futures import ThreadPoolExecutor
-import torch
-import open3d as o3d
+from concurrent.futures import ThreadPoolExecutor
 import trimesh
 
 import gradio as gr
@@ -15,41 +13,39 @@ from gradio_imageslider import ImageSlider
 import cv2
 import numpy as np
 import imageio
-# from promptda.promptda import PromptDA
-# from promptda.utils.io_wrapper import load_image, load_depth
-# from promptda.utils.depth_utils import visualize_depth, unproject_depth
-# import torch
+from promptda.promptda import PromptDA
+from promptda.utils.io_wrapper import load_image, load_depth
+from promptda.utils.depth_utils import visualize_depth, unproject_depth
 DEVICE = 'cuda' 
 # if torch.cuda.is_available(
 # ) else 'mps' if torch.backends.mps.is_available() else 'cpu'
-# model = PromptDA.from_pretrained('depth-anything/promptda_vitl').to(DEVICE).eval()
+model = PromptDA.from_pretrained('depth-anything/promptda_vitl').to(DEVICE).eval()
 # model = PromptDA.from_pretrained('depth-anything/promptda_vitl').eval()
-# thread_pool_executor = ThreadPoolExecutor(max_workers=1)
-
-# def delete_later(path: Union[str, os.PathLike], delay: int = 300):
-#     print(f"Deleting file: {path}")
-#     def _delete():
-#         try: 
-#             if os.path.isfile(path):
-#                 os.remove(path) 
-#                 print(f"Deleted file: {path}")
-#             elif os.path.isdir(path):
-#                 shutil.rmtree(path)
-#                 print(f"Deleted directory: {path}")
-#         except: 
-#             pass
-#     def _wait_and_delete():
-#         time.sleep(delay)
-#         _delete(path)
-#     thread_pool_executor.submit(_wait_and_delete)
-#     atexit.register(_delete)
+thread_pool_executor = ThreadPoolExecutor(max_workers=1)
+
+def delete_later(path: Union[str, os.PathLike], delay: int = 300):
+    print(f"Deleting file: {path}")
+    def _delete():
+        try: 
+            if os.path.isfile(path):
+                os.remove(path) 
+                print(f"Deleted file: {path}")
+            elif os.path.isdir(path):
+                shutil.rmtree(path)
+                print(f"Deleted directory: {path}")
+        except: 
+            pass
+    def _wait_and_delete():
+        time.sleep(delay)
+        _delete(path)
+    thread_pool_executor.submit(_wait_and_delete)
+    atexit.register(_delete)
 
 
-# @spaces.GPU
+@spaces.GPU
 def run_with_gpu(image, prompt_depth):
     image = image.to(DEVICE)
     prompt_depth = prompt_depth.to(DEVICE)
-    model.to(DEVICE)
     depth = model.predict(image, prompt_depth)
     depth = depth[0, 0].detach().cpu().numpy()
     return depth
@@ -58,7 +54,7 @@ def check_is_stray_scanner_app_capture(input_dir):
     assert os.path.exists(os.path.join(input_dir, 'rgb.mp4')), 'rgb.mp4 not found'
     pass
 
-@spaces.GPU
+# @spaces.GPU
 def run(input_file, resolution):
     # unzip zip file
     input_file = input_file.name
@@ -96,17 +92,19 @@ def run(input_file, resolution):
     now_max = max(color.shape[1], color.shape[0])
     scale = orig_max / now_max
     ixt[:2] = ixt[:2] / scale
-    pcd = unproject_depth(depth, ixt=ixt, color=color, ret_pcd=True)
+    points, colors = unproject_depth(depth, ixt=ixt, color=color, ret_pcd=False)
+    pcd = trimesh.PointCloud(vertices=points, colors=colors)
     ply_path = os.path.join(input_dir, f'pointcloud.ply')
-    o3d.io.write_point_cloud(ply_path, pcd)
+    pcd.export(ply_path)
+    # o3d.io.write_point_cloud(ply_path, pcd)
 
     glb_path = os.path.join(input_dir, f'pointcloud.glb')
     scene_3d = trimesh.Scene()
-    glb_colors = np.asarray(pcd.colors).astype(np.float32)
+    glb_colors = np.asarray(colors).astype(np.float32)
     glb_colors = np.concatenate([glb_colors, np.ones_like(glb_colors[:, :1])], axis=1)
     # glb_colors = (np.asarray(pcd.colors) * 255).astype(np.uint8)
     pcd_data = trimesh.PointCloud(
-        vertices=np.asarray(pcd.points) * np.array([[1, -1, -1]]),
+        vertices=np.asarray(points) * np.array([[1, -1, -1]]),
         colors=glb_colors.astype(np.float64),
     )
     scene_3d.add_geometry(pcd_data)

promptda/utils/depth_utils.py

@@ -1,6 +1,5 @@
 import numpy as np
 import matplotlib
-import open3d as o3d
 
 def visualize_depth(depth: np.ndarray, 
                     depth_min=None, 
@@ -72,6 +71,7 @@ def unproject_depth(depth,
             color = color.astype(np.float32) / 255.
         if ret_pcd:
             points = pcd
+            import open3d as o3d
             pcd = o3d.geometry.PointCloud()
             pcd.points = o3d.utility.Vector3dVector(points[:, :3][new_mask])
             pcd.colors = o3d.utility.Vector3dVector(color.reshape(-1, 3)[mask][new_mask])
@@ -79,6 +79,7 @@ def unproject_depth(depth,
             return pcd[:, :3][new_mask], color.reshape(-1, 3)[mask][new_mask]
     else:
         if ret_pcd:
+            import open3d as o3d
             points = pcd
             pcd = o3d.geometry.PointCloud()
             pcd.points = o3d.utility.Vector3dVector(pcd[:, :3][new_mask])

requirements.txt

@@ -19,7 +19,7 @@ opencv-python==4.9.0.80
 scipy
 matplotlib
 h5py
-open3d
+# open3d
 
 # app.py
 gradio==4.44.1