Spaces:
Running
on
Zero
Running
on
Zero
haotongl
commited on
Commit
•
ae88fe1
1
Parent(s):
389b85f
inital version
Browse files- app.py +36 -38
- promptda/utils/depth_utils.py +2 -1
- requirements.txt +1 -1
app.py
CHANGED
@@ -1,13 +1,11 @@
|
|
1 |
import os
|
2 |
-
|
3 |
-
|
4 |
from pathlib import Path
|
5 |
-
|
6 |
-
|
7 |
import spaces
|
8 |
-
|
9 |
-
import torch
|
10 |
-
import open3d as o3d
|
11 |
import trimesh
|
12 |
|
13 |
import gradio as gr
|
@@ -15,41 +13,39 @@ from gradio_imageslider import ImageSlider
|
|
15 |
import cv2
|
16 |
import numpy as np
|
17 |
import imageio
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
# import torch
|
22 |
DEVICE = 'cuda'
|
23 |
# if torch.cuda.is_available(
|
24 |
# ) else 'mps' if torch.backends.mps.is_available() else 'cpu'
|
25 |
-
|
26 |
# model = PromptDA.from_pretrained('depth-anything/promptda_vitl').eval()
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
|
47 |
|
48 |
-
|
49 |
def run_with_gpu(image, prompt_depth):
|
50 |
image = image.to(DEVICE)
|
51 |
prompt_depth = prompt_depth.to(DEVICE)
|
52 |
-
model.to(DEVICE)
|
53 |
depth = model.predict(image, prompt_depth)
|
54 |
depth = depth[0, 0].detach().cpu().numpy()
|
55 |
return depth
|
@@ -58,7 +54,7 @@ def check_is_stray_scanner_app_capture(input_dir):
|
|
58 |
assert os.path.exists(os.path.join(input_dir, 'rgb.mp4')), 'rgb.mp4 not found'
|
59 |
pass
|
60 |
|
61 |
-
@spaces.GPU
|
62 |
def run(input_file, resolution):
|
63 |
# unzip zip file
|
64 |
input_file = input_file.name
|
@@ -96,17 +92,19 @@ def run(input_file, resolution):
|
|
96 |
now_max = max(color.shape[1], color.shape[0])
|
97 |
scale = orig_max / now_max
|
98 |
ixt[:2] = ixt[:2] / scale
|
99 |
-
|
|
|
100 |
ply_path = os.path.join(input_dir, f'pointcloud.ply')
|
101 |
-
|
|
|
102 |
|
103 |
glb_path = os.path.join(input_dir, f'pointcloud.glb')
|
104 |
scene_3d = trimesh.Scene()
|
105 |
-
glb_colors = np.asarray(
|
106 |
glb_colors = np.concatenate([glb_colors, np.ones_like(glb_colors[:, :1])], axis=1)
|
107 |
# glb_colors = (np.asarray(pcd.colors) * 255).astype(np.uint8)
|
108 |
pcd_data = trimesh.PointCloud(
|
109 |
-
vertices=np.asarray(
|
110 |
colors=glb_colors.astype(np.float64),
|
111 |
)
|
112 |
scene_3d.add_geometry(pcd_data)
|
|
|
1 |
import os
|
2 |
+
import time
|
3 |
+
import shutil
|
4 |
from pathlib import Path
|
5 |
+
from typing import Union
|
6 |
+
import atexit
|
7 |
import spaces
|
8 |
+
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
|
9 |
import trimesh
|
10 |
|
11 |
import gradio as gr
|
|
|
13 |
import cv2
|
14 |
import numpy as np
|
15 |
import imageio
|
16 |
+
from promptda.promptda import PromptDA
|
17 |
+
from promptda.utils.io_wrapper import load_image, load_depth
|
18 |
+
from promptda.utils.depth_utils import visualize_depth, unproject_depth
|
|
|
19 |
DEVICE = 'cuda'
|
20 |
# if torch.cuda.is_available(
|
21 |
# ) else 'mps' if torch.backends.mps.is_available() else 'cpu'
|
22 |
+
model = PromptDA.from_pretrained('depth-anything/promptda_vitl').to(DEVICE).eval()
|
23 |
# model = PromptDA.from_pretrained('depth-anything/promptda_vitl').eval()
|
24 |
+
thread_pool_executor = ThreadPoolExecutor(max_workers=1)
|
25 |
+
|
26 |
+
def delete_later(path: Union[str, os.PathLike], delay: int = 300):
|
27 |
+
print(f"Deleting file: {path}")
|
28 |
+
def _delete():
|
29 |
+
try:
|
30 |
+
if os.path.isfile(path):
|
31 |
+
os.remove(path)
|
32 |
+
print(f"Deleted file: {path}")
|
33 |
+
elif os.path.isdir(path):
|
34 |
+
shutil.rmtree(path)
|
35 |
+
print(f"Deleted directory: {path}")
|
36 |
+
except:
|
37 |
+
pass
|
38 |
+
def _wait_and_delete():
|
39 |
+
time.sleep(delay)
|
40 |
+
_delete(path)
|
41 |
+
thread_pool_executor.submit(_wait_and_delete)
|
42 |
+
atexit.register(_delete)
|
43 |
|
44 |
|
45 |
+
@spaces.GPU
|
46 |
def run_with_gpu(image, prompt_depth):
|
47 |
image = image.to(DEVICE)
|
48 |
prompt_depth = prompt_depth.to(DEVICE)
|
|
|
49 |
depth = model.predict(image, prompt_depth)
|
50 |
depth = depth[0, 0].detach().cpu().numpy()
|
51 |
return depth
|
|
|
54 |
assert os.path.exists(os.path.join(input_dir, 'rgb.mp4')), 'rgb.mp4 not found'
|
55 |
pass
|
56 |
|
57 |
+
# @spaces.GPU
|
58 |
def run(input_file, resolution):
|
59 |
# unzip zip file
|
60 |
input_file = input_file.name
|
|
|
92 |
now_max = max(color.shape[1], color.shape[0])
|
93 |
scale = orig_max / now_max
|
94 |
ixt[:2] = ixt[:2] / scale
|
95 |
+
points, colors = unproject_depth(depth, ixt=ixt, color=color, ret_pcd=False)
|
96 |
+
pcd = trimesh.PointCloud(vertices=points, colors=colors)
|
97 |
ply_path = os.path.join(input_dir, f'pointcloud.ply')
|
98 |
+
pcd.export(ply_path)
|
99 |
+
# o3d.io.write_point_cloud(ply_path, pcd)
|
100 |
|
101 |
glb_path = os.path.join(input_dir, f'pointcloud.glb')
|
102 |
scene_3d = trimesh.Scene()
|
103 |
+
glb_colors = np.asarray(colors).astype(np.float32)
|
104 |
glb_colors = np.concatenate([glb_colors, np.ones_like(glb_colors[:, :1])], axis=1)
|
105 |
# glb_colors = (np.asarray(pcd.colors) * 255).astype(np.uint8)
|
106 |
pcd_data = trimesh.PointCloud(
|
107 |
+
vertices=np.asarray(points) * np.array([[1, -1, -1]]),
|
108 |
colors=glb_colors.astype(np.float64),
|
109 |
)
|
110 |
scene_3d.add_geometry(pcd_data)
|
promptda/utils/depth_utils.py
CHANGED
@@ -1,6 +1,5 @@
|
|
1 |
import numpy as np
|
2 |
import matplotlib
|
3 |
-
import open3d as o3d
|
4 |
|
5 |
def visualize_depth(depth: np.ndarray,
|
6 |
depth_min=None,
|
@@ -72,6 +71,7 @@ def unproject_depth(depth,
|
|
72 |
color = color.astype(np.float32) / 255.
|
73 |
if ret_pcd:
|
74 |
points = pcd
|
|
|
75 |
pcd = o3d.geometry.PointCloud()
|
76 |
pcd.points = o3d.utility.Vector3dVector(points[:, :3][new_mask])
|
77 |
pcd.colors = o3d.utility.Vector3dVector(color.reshape(-1, 3)[mask][new_mask])
|
@@ -79,6 +79,7 @@ def unproject_depth(depth,
|
|
79 |
return pcd[:, :3][new_mask], color.reshape(-1, 3)[mask][new_mask]
|
80 |
else:
|
81 |
if ret_pcd:
|
|
|
82 |
points = pcd
|
83 |
pcd = o3d.geometry.PointCloud()
|
84 |
pcd.points = o3d.utility.Vector3dVector(pcd[:, :3][new_mask])
|
|
|
1 |
import numpy as np
|
2 |
import matplotlib
|
|
|
3 |
|
4 |
def visualize_depth(depth: np.ndarray,
|
5 |
depth_min=None,
|
|
|
71 |
color = color.astype(np.float32) / 255.
|
72 |
if ret_pcd:
|
73 |
points = pcd
|
74 |
+
import open3d as o3d
|
75 |
pcd = o3d.geometry.PointCloud()
|
76 |
pcd.points = o3d.utility.Vector3dVector(points[:, :3][new_mask])
|
77 |
pcd.colors = o3d.utility.Vector3dVector(color.reshape(-1, 3)[mask][new_mask])
|
|
|
79 |
return pcd[:, :3][new_mask], color.reshape(-1, 3)[mask][new_mask]
|
80 |
else:
|
81 |
if ret_pcd:
|
82 |
+
import open3d as o3d
|
83 |
points = pcd
|
84 |
pcd = o3d.geometry.PointCloud()
|
85 |
pcd.points = o3d.utility.Vector3dVector(pcd[:, :3][new_mask])
|
requirements.txt
CHANGED
@@ -19,7 +19,7 @@ opencv-python==4.9.0.80
|
|
19 |
scipy
|
20 |
matplotlib
|
21 |
h5py
|
22 |
-
open3d
|
23 |
|
24 |
# app.py
|
25 |
gradio==4.44.1
|
|
|
19 |
scipy
|
20 |
matplotlib
|
21 |
h5py
|
22 |
+
# open3d
|
23 |
|
24 |
# app.py
|
25 |
gradio==4.44.1
|