hololens/stable-diffusion-webui-depthmap-script
1
1import torch
2import numpy as np
3from torchsparse import SparseTensor
4from torchsparse.utils import sparse_collate_fn, sparse_quantize
5from plyfile import PlyData, PlyElement
6import os
7
8def init_image_coor(height, width, u0=None, v0=None):
9 u0 = width / 2.0 if u0 is None else u0
10 v0 = height / 2.0 if v0 is None else v0
11
12 x_row = np.arange(0, width)
13 x = np.tile(x_row, (height, 1))
14 x = x.astype(np.float32)
15 u_u0 = x - u0
16
17 y_col = np.arange(0, height)
18 y = np.tile(y_col, (width, 1)).T
19 y = y.astype(np.float32)
20 v_v0 = y - v0
21 return u_u0, v_v0
22
23def depth_to_pcd(depth, u_u0, v_v0, f, invalid_value=0):
24 mask_invalid = depth <= invalid_value
25 depth[mask_invalid] = 0.0
26 x = u_u0 / f * depth
27 y = v_v0 / f * depth
28 z = depth
29 pcd = np.stack([x, y, z], axis=2)
30 return pcd, ~mask_invalid
31
32def pcd_to_sparsetensor(pcd, mask_valid, voxel_size=0.01, num_points=100000):
33 pcd_valid = pcd[mask_valid]
34 block_ = pcd_valid
35 block = np.zeros_like(block_)
36 block[:, :3] = block_[:, :3]
37
38 pc_ = np.round(block_[:, :3] / voxel_size)
39 pc_ -= pc_.min(0, keepdims=1)
40 feat_ = block
41
42 # transfer point cloud to voxels
43 inds = sparse_quantize(pc_,
44 feat_,
45 return_index=True,
46 return_invs=False)
47 if len(inds) > num_points:
48 inds = np.random.choice(inds, num_points, replace=False)
49
50 pc = pc_[inds]
51 feat = feat_[inds]
52 lidar = SparseTensor(feat, pc)
53 feed_dict = [{'lidar': lidar}]
54 inputs = sparse_collate_fn(feed_dict)
55 return inputs
56
57def pcd_uv_to_sparsetensor(pcd, u_u0, v_v0, mask_valid, f= 500.0, voxel_size=0.01, mask_side=None, num_points=100000):
58 if mask_side is not None:
59 mask_valid = mask_valid & mask_side
60 pcd_valid = pcd[mask_valid]
61 u_u0_valid = u_u0[mask_valid][:, np.newaxis] / f
62 v_v0_valid = v_v0[mask_valid][:, np.newaxis] / f
63
64 block_ = np.concatenate([pcd_valid, u_u0_valid, v_v0_valid], axis=1)
65 block = np.zeros_like(block_)
66 block[:, :] = block_[:, :]
67
68
69 pc_ = np.round(block_[:, :3] / voxel_size)
70 pc_ -= pc_.min(0, keepdims=1)
71 feat_ = block
72
73 # transfer point cloud to voxels
74 inds = sparse_quantize(pc_,
75 feat_,
76 return_index=True,
77 return_invs=False)
78 if len(inds) > num_points:
79 inds = np.random.choice(inds, num_points, replace=False)
80
81 pc = pc_[inds]
82 feat = feat_[inds]
83 lidar = SparseTensor(feat, pc)
84 feed_dict = [{'lidar': lidar}]
85 inputs = sparse_collate_fn(feed_dict)
86 return inputs
87
88
89def refine_focal_one_step(depth, focal, model, u0, v0):
90 # reconstruct PCD from depth
91 u_u0, v_v0 = init_image_coor(depth.shape[0], depth.shape[1], u0=u0, v0=v0)
92 pcd, mask_valid = depth_to_pcd(depth, u_u0, v_v0, f=focal, invalid_value=0)
93 # input for the voxelnet
94 feed_dict = pcd_uv_to_sparsetensor(pcd, u_u0, v_v0, mask_valid, f=focal, voxel_size=0.005, mask_side=None)
95 inputs = feed_dict['lidar'].cuda()
96
97 outputs = model(inputs)
98 return outputs
99
100def refine_shift_one_step(depth_wshift, model, focal, u0, v0):
101 # reconstruct PCD from depth
102 u_u0, v_v0 = init_image_coor(depth_wshift.shape[0], depth_wshift.shape[1], u0=u0, v0=v0)
103 pcd_wshift, mask_valid = depth_to_pcd(depth_wshift, u_u0, v_v0, f=focal, invalid_value=0)
104 # input for the voxelnet
105 feed_dict = pcd_to_sparsetensor(pcd_wshift, mask_valid, voxel_size=0.01)
106 inputs = feed_dict['lidar'].cuda()
107
108 outputs = model(inputs)
109 return outputs
110
111def refine_focal(depth, focal, model, u0, v0):
112 last_scale = 1
113 focal_tmp = np.copy(focal)
114 for i in range(1):
115 scale = refine_focal_one_step(depth, focal_tmp, model, u0, v0)
116 focal_tmp = focal_tmp / scale.item()
117 last_scale = last_scale * scale
118 return torch.tensor([[last_scale]])
119
120def refine_shift(depth_wshift, model, focal, u0, v0):
121 depth_wshift_tmp = np.copy(depth_wshift)
122 last_shift = 0
123 for i in range(1):
124 shift = refine_shift_one_step(depth_wshift_tmp, model, focal, u0, v0)
125 shift = shift if shift.item() < 0.7 else torch.tensor([[0.7]])
126 depth_wshift_tmp -= shift.item()
127 last_shift += shift.item()
128 return torch.tensor([[last_shift]])
129
130def reconstruct_3D(depth, f):
131 """
132 Reconstruct depth to 3D pointcloud with the provided focal length.
133 Return:
134 pcd: N X 3 array, point cloud
135 """
136 cu = depth.shape[1] / 2
137 cv = depth.shape[0] / 2
138 width = depth.shape[1]
139 height = depth.shape[0]
140 row = np.arange(0, width, 1)
141 u = np.array([row for i in np.arange(height)])
142 col = np.arange(0, height, 1)
143 v = np.array([col for i in np.arange(width)])
144 v = v.transpose(1, 0)
145
146 if f > 1e5:
147 print('Infinit focal length!!!')
148 x = u - cu
149 y = v - cv
150 z = depth / depth.max() * x.max()
151 else:
152 x = (u - cu) * depth / f
153 y = (v - cv) * depth / f
154 z = depth
155
156 x = np.reshape(x, (width * height, 1)).astype(float)
157 y = np.reshape(y, (width * height, 1)).astype(float)
158 z = np.reshape(z, (width * height, 1)).astype(float)
159 pcd = np.concatenate((x, y, z), axis=1)
160 pcd = pcd.astype(int)
161 return pcd
162
163def save_point_cloud(pcd, rgb, filename, binary=True):
164 """Save an RGB point cloud as a PLY file.
165
166 :paras
167 @pcd: Nx3 matrix, the XYZ coordinates
168 @rgb: NX3 matrix, the rgb colors for each 3D point
169 """
170 assert pcd.shape[0] == rgb.shape[0]
171
172 if rgb is None:
173 gray_concat = np.tile(np.array([128], dtype=np.uint8), (pcd.shape[0], 3))
174 points_3d = np.hstack((pcd, gray_concat))
175 else:
176 points_3d = np.hstack((pcd, rgb))
177 python_types = (float, float, float, int, int, int)
178 npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),
179 ('blue', 'u1')]
180 if binary is True:
181 # Format into NumPy structured array
182 vertices = []
183 for row_idx in range(points_3d.shape[0]):
184 cur_point = points_3d[row_idx]
185 vertices.append(tuple(dtype(point) for dtype, point in zip(python_types, cur_point)))
186 vertices_array = np.array(vertices, dtype=npy_types)
187 el = PlyElement.describe(vertices_array, 'vertex')
188
189 # Write
190 PlyData([el]).write(filename)
191 else:
192 x = np.squeeze(points_3d[:, 0])
193 y = np.squeeze(points_3d[:, 1])
194 z = np.squeeze(points_3d[:, 2])
195 r = np.squeeze(points_3d[:, 3])
196 g = np.squeeze(points_3d[:, 4])
197 b = np.squeeze(points_3d[:, 5])
198
199 ply_head = 'ply\n' \
200 'format ascii 1.0\n' \
201 'element vertex %d\n' \
202 'property float x\n' \
203 'property float y\n' \
204 'property float z\n' \
205 'property uchar red\n' \
206 'property uchar green\n' \
207 'property uchar blue\n' \
208 'end_header' % r.shape[0]
209 # ---- Save ply data to disk
210 np.savetxt(filename, np.column_stack((x, y, z, r, g, b)), fmt="%d %d %d %d %d %d", header=ply_head, comments='')
211
212def reconstruct_depth(depth, rgb, dir, pcd_name, focal):
213 """
214 para disp: disparity, [h, w]
215 para rgb: rgb image, [h, w, 3], in rgb format
216 """
217 rgb = np.squeeze(rgb)
218 depth = np.squeeze(depth)
219
220 mask = depth < 1e-8
221 depth[mask] = 0
222 depth = depth / depth.max() * 10000
223
224 pcd = reconstruct_3D(depth, f=focal)
225 rgb_n = np.reshape(rgb, (-1, 3))
226 save_point_cloud(pcd, rgb_n, os.path.join(dir, pcd_name + '.ply'))
227
228
229def recover_metric_depth(pred, gt):
230 if type(pred).__module__ == torch.__name__:
231 pred = pred.cpu().numpy()
232 if type(gt).__module__ == torch.__name__:
233 gt = gt.cpu().numpy()
234 gt = gt.squeeze()
235 pred = pred.squeeze()
236 mask = (gt > 1e-8) & (pred > 1e-8)
237
238 gt_mask = gt[mask]
239 pred_mask = pred[mask]
240 a, b = np.polyfit(pred_mask, gt_mask, deg=1)
241 pred_metric = a * pred + b
242 return pred_metric
243