hololens/stable-diffusion-webui-depthmap-script
1
1import os
2import numpy as np
3try:
4 import cynetworkx as netx
5except ImportError:
6 import networkx as netx
7
8import json
9import scipy.misc as misc
10#import OpenEXR
11import scipy.signal as signal
12import matplotlib.pyplot as plt
13import cv2
14import scipy.misc as misc
15from skimage import io
16from functools import partial
17from vispy import scene, io
18from vispy.scene import visuals
19from functools import reduce
20# from moviepy.editor import ImageSequenceClip
21import scipy.misc as misc
22from vispy.visuals.filters import Alpha
23import cv2
24from skimage.transform import resize
25import copy
26import torch
27import os
28from inpaint.utils import refine_depth_around_edge, smooth_cntsyn_gap
29from inpaint.utils import require_depth_edge, filter_irrelevant_edge_new, open_small_mask
30from skimage.feature import canny
31from scipy import ndimage
32import time
33import transforms3d
34
35def relabel_node(mesh, nodes, cur_node, new_node):
36 if cur_node == new_node:
37 return mesh
38 mesh.add_node(new_node)
39 for key, value in nodes[cur_node].items():
40 nodes[new_node][key] = value
41 for ne in mesh.neighbors(cur_node):
42 mesh.add_edge(new_node, ne)
43 mesh.remove_node(cur_node)
44
45 return mesh
46
47def filter_edge(mesh, edge_ccs, config, invalid=False):
48 context_ccs = [set() for _ in edge_ccs]
49 mesh_nodes = mesh.nodes
50 for edge_id, edge_cc in enumerate(edge_ccs):
51 if config['context_thickness'] == 0:
52 continue
53 edge_group = {}
54 for edge_node in edge_cc:
55 far_nodes = mesh_nodes[edge_node].get('far')
56 if far_nodes is None:
57 continue
58 for far_node in far_nodes:
59 context_ccs[edge_id].add(far_node)
60 if mesh_nodes[far_node].get('edge_id') is not None:
61 if edge_group.get(mesh_nodes[far_node]['edge_id']) is None:
62 edge_group[mesh_nodes[far_node]['edge_id']] = set()
63 edge_group[mesh_nodes[far_node]['edge_id']].add(far_node)
64 if len(edge_cc) > 2:
65 for edge_key in [*edge_group.keys()]:
66 if len(edge_group[edge_key]) == 1:
67 context_ccs[edge_id].remove([*edge_group[edge_key]][0])
68 valid_edge_ccs = []
69 for xidx, yy in enumerate(edge_ccs):
70 if invalid is not True and len(context_ccs[xidx]) > 0:
71 # if len(context_ccs[xidx]) > 0:
72 valid_edge_ccs.append(yy)
73 elif invalid is True and len(context_ccs[xidx]) == 0:
74 valid_edge_ccs.append(yy)
75 else:
76 valid_edge_ccs.append(set())
77 # valid_edge_ccs = [yy for xidx, yy in enumerate(edge_ccs) if len(context_ccs[xidx]) > 0]
78
79 return valid_edge_ccs
80
81def extrapolate(global_mesh,
82 info_on_pix,
83 image,
84 depth,
85 other_edge_with_id,
86 edge_map,
87 edge_ccs,
88 depth_edge_model,
89 depth_feat_model,
90 rgb_feat_model,
91 config,
92 direc='right-up'):
93 h_off, w_off = global_mesh.graph['hoffset'], global_mesh.graph['woffset']
94 noext_H, noext_W = global_mesh.graph['noext_H'], global_mesh.graph['noext_W']
95
96 if "up" in direc.lower() and "-" not in direc.lower():
97 all_anchor = [0, h_off + config['context_thickness'], w_off, w_off + noext_W]
98 global_shift = [all_anchor[0], all_anchor[2]]
99 mask_anchor = [0, h_off, w_off, w_off + noext_W]
100 context_anchor = [h_off, h_off + config['context_thickness'], w_off, w_off + noext_W]
101 valid_line_anchor = [h_off, h_off + 1, w_off, w_off + noext_W]
102 valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
103 min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
104 elif "down" in direc.lower() and "-" not in direc.lower():
105 all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, w_off, w_off + noext_W]
106 global_shift = [all_anchor[0], all_anchor[2]]
107 mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off, w_off + noext_W]
108 context_anchor = [h_off + noext_H - config['context_thickness'], h_off + noext_H, w_off, w_off + noext_W]
109 valid_line_anchor = [h_off + noext_H - 1, h_off + noext_H, w_off, w_off + noext_W]
110 valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
111 min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
112 elif "left" in direc.lower() and "-" not in direc.lower():
113 all_anchor = [h_off, h_off + noext_H, 0, w_off + config['context_thickness']]
114 global_shift = [all_anchor[0], all_anchor[2]]
115 mask_anchor = [h_off, h_off + noext_H, 0, w_off]
116 context_anchor = [h_off, h_off + noext_H, w_off, w_off + config['context_thickness']]
117 valid_line_anchor = [h_off, h_off + noext_H, w_off, w_off + 1]
118 valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
119 min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
120 elif "right" in direc.lower() and "-" not in direc.lower():
121 all_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
122 global_shift = [all_anchor[0], all_anchor[2]]
123 mask_anchor = [h_off, h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
124 context_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], w_off + noext_W]
125 valid_line_anchor = [h_off, h_off + noext_H, w_off + noext_W - 1, w_off + noext_W]
126 valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
127 min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
128 elif "left" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
129 all_anchor = [0, h_off + config['context_thickness'], 0, w_off + config['context_thickness']]
130 global_shift = [all_anchor[0], all_anchor[2]]
131 mask_anchor = [0, h_off, 0, w_off]
132 context_anchor = "inv-mask"
133 valid_line_anchor = None
134 valid_anchor = all_anchor
135 elif "left" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
136 all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, 0, w_off + config['context_thickness']]
137 global_shift = [all_anchor[0], all_anchor[2]]
138 mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, 0, w_off]
139 context_anchor = "inv-mask"
140 valid_line_anchor = None
141 valid_anchor = all_anchor
142 elif "right" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
143 all_anchor = [0, h_off + config['context_thickness'], w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
144 global_shift = [all_anchor[0], all_anchor[2]]
145 mask_anchor = [0, h_off, w_off + noext_W, 2 * w_off + noext_W]
146 context_anchor = "inv-mask"
147 valid_line_anchor = None
148 valid_anchor = all_anchor
149 elif "right" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
150 all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
151 global_shift = [all_anchor[0], all_anchor[2]]
152 mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
153 context_anchor = "inv-mask"
154 valid_line_anchor = None
155 valid_anchor = all_anchor
156
157 global_mask = np.zeros_like(depth)
158 global_mask[mask_anchor[0]:mask_anchor[1],mask_anchor[2]:mask_anchor[3]] = 1
159 mask = global_mask[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1
160 context = 1 - mask
161 global_context = np.zeros_like(depth)
162 global_context[all_anchor[0]:all_anchor[1],all_anchor[2]:all_anchor[3]] = context
163 # context = global_context[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1
164
165
166
167 valid_area = mask + context
168 input_rgb = image[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] / 255. * context[..., None]
169 input_depth = depth[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * context
170 log_depth = np.log(input_depth + 1e-8)
171 log_depth[mask > 0] = 0
172 input_mean_depth = np.mean(log_depth[context > 0])
173 input_zero_mean_depth = (log_depth - input_mean_depth) * context
174 input_disp = 1./np.abs(input_depth)
175 input_disp[mask > 0] = 0
176 input_disp = input_disp / input_disp.max()
177 valid_line = np.zeros_like(depth)
178 if valid_line_anchor is not None:
179 valid_line[valid_line_anchor[0]:valid_line_anchor[1], valid_line_anchor[2]:valid_line_anchor[3]] = 1
180 valid_line = valid_line[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
181 # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(global_context * 1 + global_mask * 2); ax2.imshow(image); plt.show()
182 # f, ((ax1, ax2, ax3)) = plt.subplots(1, 3, sharex=True, sharey=True); ax1.imshow(context * 1 + mask * 2); ax2.imshow(input_rgb); ax3.imshow(valid_line); plt.show()
183 # import pdb; pdb.set_trace()
184 # return
185 input_edge_map = edge_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] * context
186 input_other_edge_with_id = other_edge_with_id[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
187 end_depth_maps = ((valid_line * input_edge_map) > 0) * input_depth
188
189
190 if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
191 device = config["gpu_ids"]
192 else:
193 device = "cpu"
194
195 valid_edge_ids = sorted(list(input_other_edge_with_id[(valid_line * input_edge_map) > 0]))
196 valid_edge_ids = valid_edge_ids[1:] if (len(valid_edge_ids) > 0 and valid_edge_ids[0] == -1) else valid_edge_ids
197 edge = reduce(lambda x, y: (x + (input_other_edge_with_id == y).astype(np.uint8)).clip(0, 1), [np.zeros_like(mask)] + list(valid_edge_ids))
198 t_edge = torch.FloatTensor(edge).to(device)[None, None, ...]
199 t_rgb = torch.FloatTensor(input_rgb).to(device).permute(2,0,1).unsqueeze(0)
200 t_mask = torch.FloatTensor(mask).to(device)[None, None, ...]
201 t_context = torch.FloatTensor(context).to(device)[None, None, ...]
202 t_disp = torch.FloatTensor(input_disp).to(device)[None, None, ...]
203 t_depth_zero_mean_depth = torch.FloatTensor(input_zero_mean_depth).to(device)[None, None, ...]
204
205 depth_edge_output = depth_edge_model.forward_3P(t_mask, t_context, t_rgb, t_disp, t_edge, unit_length=128,
206 cuda=device)
207 t_output_edge = (depth_edge_output> config['ext_edge_threshold']).float() * t_mask + t_edge
208 output_raw_edge = t_output_edge.data.cpu().numpy().squeeze()
209 # import pdb; pdb.set_trace()
210 mesh = netx.Graph()
211 hxs, hys = np.where(output_raw_edge * mask > 0)
212 valid_map = mask + context
213 for hx, hy in zip(hxs, hys):
214 node = (hx, hy)
215 mesh.add_node((hx, hy))
216 eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
217 (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\
218 if 0 <= ne[0] < output_raw_edge.shape[0] and 0 <= ne[1] < output_raw_edge.shape[1] and 0 < output_raw_edge[ne[0], ne[1]]]
219 for ne in eight_nes:
220 mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy))
221 if end_depth_maps[ne[0], ne[1]] != 0:
222 mesh.nodes[ne[0], ne[1]]['cnt'] = True
223 mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
224 ccs = [*netx.connected_components(mesh)]
225 end_pts = []
226 for cc in ccs:
227 end_pts.append(set())
228 for node in cc:
229 if mesh.nodes[node].get('cnt') is not None:
230 end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth']))
231 fpath_map = np.zeros_like(output_raw_edge) - 1
232 npath_map = np.zeros_like(output_raw_edge) - 1
233 for end_pt, cc in zip(end_pts, ccs):
234 sorted_end_pt = []
235 if len(end_pt) >= 2:
236 continue
237 if len(end_pt) == 0:
238 continue
239 if len(end_pt) == 1:
240 sub_mesh = mesh.subgraph(list(cc)).copy()
241 pnodes = netx.periphery(sub_mesh)
242 ends = [*end_pt]
243 edge_id = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])]['edge_id']
244 pnodes = sorted(pnodes,
245 key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
246 reverse=True)[0]
247 npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
248 for np_node in npath:
249 npath_map[np_node[0], np_node[1]] = edge_id
250 fpath = []
251 if global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far') is None:
252 print("None far")
253 import pdb; pdb.set_trace()
254 else:
255 fnodes = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far')
256 fnodes = [(xx[0] - all_anchor[0], xx[1] - all_anchor[2], xx[2]) for xx in fnodes]
257 dmask = mask + 0
258 did = 0
259 while True:
260 did += 1
261 dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
262 if did > 3:
263 break
264 # ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0)]
265 ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0)]
266 if len(ffnode) > 0:
267 fnode = ffnode[0]
268 break
269 if len(ffnode) == 0:
270 continue
271 fpath.append((fnode[0], fnode[1]))
272 for step in range(0, len(npath) - 1):
273 parr = (npath[step + 1][0] - npath[step][0], npath[step + 1][1] - npath[step][1])
274 new_loc = (fpath[-1][0] + parr[0], fpath[-1][1] + parr[1])
275 new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
276 (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
277 if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
278 if np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != -4:
279 break
280 if npath_map[new_loc[0], new_loc[1]] != -1:
281 if npath_map[new_loc[0], new_loc[1]] != edge_id:
282 break
283 else:
284 continue
285 if valid_area[new_loc[0], new_loc[1]] == 0:
286 break
287 new_loc_nes_eight = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
288 (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1),
289 (new_loc[0] + 1, new_loc[1] + 1), (new_loc[0] + 1, new_loc[1] - 1),
290 (new_loc[0] - 1, new_loc[1] - 1), (new_loc[0] - 1, new_loc[1] + 1)]\
291 if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
292 if np.sum([int(npath_map[nlne[0], nlne[1]] == edge_id) for nlne in new_loc_nes_eight]) == 0:
293 break
294 fpath.append((fpath[-1][0] + parr[0], fpath[-1][1] + parr[1]))
295 if step != len(npath) - 2:
296 for xx in npath[step+1:]:
297 if npath_map[xx[0], xx[1]] == edge_id:
298 npath_map[xx[0], xx[1]] = -1
299 if len(fpath) > 0:
300 for fp_node in fpath:
301 fpath_map[fp_node[0], fp_node[1]] = edge_id
302 # import pdb; pdb.set_trace()
303 far_edge = (fpath_map > -1).astype(np.uint8)
304 update_edge = (npath_map > -1) * mask + edge
305 t_update_edge = torch.FloatTensor(update_edge).to(device)[None, None, ...]
306 depth_output = depth_feat_model.forward_3P(t_mask, t_context, t_depth_zero_mean_depth, t_update_edge, unit_length=128,
307 cuda=device)
308 depth_output = depth_output.cpu().data.numpy().squeeze()
309 depth_output = np.exp(depth_output + input_mean_depth) * mask # + input_depth * context
310 # if "right" in direc.lower() and "-" not in direc.lower():
311 # plt.imshow(depth_output); plt.show()
312 # import pdb; pdb.set_trace()
313 # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
314 for near_id in np.unique(npath_map[npath_map > -1]):
315 depth_output = refine_depth_around_edge(depth_output.copy(),
316 (fpath_map == near_id).astype(np.uint8) * mask, # far_edge_map_in_mask,
317 (fpath_map == near_id).astype(np.uint8), # far_edge_map,
318 (npath_map == near_id).astype(np.uint8) * mask,
319 mask.copy(),
320 np.zeros_like(mask),
321 config)
322 # if "right" in direc.lower() and "-" not in direc.lower():
323 # plt.imshow(depth_output); plt.show()
324 # import pdb; pdb.set_trace()
325 # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
326 rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128,
327 cuda=device)
328
329 # rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128, cuda=config['gpu_ids'])
330 if config.get('gray_image') is True:
331 rgb_output = rgb_output.mean(1, keepdim=True).repeat((1,3,1,1))
332 rgb_output = ((rgb_output.squeeze().data.cpu().permute(1,2,0).numpy() * mask[..., None] + input_rgb) * 255).astype(np.uint8)
333 image[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]][mask > 0] = rgb_output[mask > 0] # np.array([255,0,0]) # rgb_output[mask > 0]
334 depth[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]][mask > 0] = depth_output[mask > 0]
335 # nxs, nys = np.where(mask > -1)
336 # for nx, ny in zip(nxs, nys):
337 # info_on_pix[(nx, ny)][0]['color'] = rgb_output[]
338
339
340 nxs, nys = np.where((npath_map > -1))
341 for nx, ny in zip(nxs, nys):
342 n_id = npath_map[nx, ny]
343 four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
344 if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
345 for nex, ney in four_nes:
346 if fpath_map[nex, ney] == n_id:
347 na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
348 (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
349 if global_mesh.has_edge(na, nb):
350 global_mesh.remove_edge(na, nb)
351 nxs, nys = np.where((fpath_map > -1))
352 for nx, ny in zip(nxs, nys):
353 n_id = fpath_map[nx, ny]
354 four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
355 if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
356 for nex, ney in four_nes:
357 if npath_map[nex, ney] == n_id:
358 na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
359 (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
360 if global_mesh.has_edge(na, nb):
361 global_mesh.remove_edge(na, nb)
362 nxs, nys = np.where(mask > 0)
363 for x, y in zip(nxs, nys):
364 x = x + all_anchor[0]
365 y = y + all_anchor[2]
366 cur_node = (x, y, 0)
367 new_node = (x, y, -abs(depth[x, y]))
368 disp = 1. / -abs(depth[x, y])
369 mapping_dict = {cur_node: new_node}
370 info_on_pix, global_mesh = update_info(mapping_dict, info_on_pix, global_mesh)
371 global_mesh.nodes[new_node]['color'] = image[x, y]
372 global_mesh.nodes[new_node]['old_color'] = image[x, y]
373 global_mesh.nodes[new_node]['disp'] = disp
374 info_on_pix[(x, y)][0]['depth'] = -abs(depth[x, y])
375 info_on_pix[(x, y)][0]['disp'] = disp
376 info_on_pix[(x, y)][0]['color'] = image[x, y]
377
378
379 nxs, nys = np.where((npath_map > -1))
380 for nx, ny in zip(nxs, nys):
381 self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
382 if global_mesh.has_node(self_node) is False:
383 break
384 n_id = int(round(npath_map[nx, ny]))
385 four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
386 if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
387 for nex, ney in four_nes:
388 ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
389 if global_mesh.has_node(ne_node) is False:
390 continue
391 if fpath_map[nex, ney] == n_id:
392 if global_mesh.nodes[self_node].get('edge_id') is None:
393 global_mesh.nodes[self_node]['edge_id'] = n_id
394 edge_ccs[n_id].add(self_node)
395 info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = n_id
396 if global_mesh.has_edge(self_node, ne_node) is True:
397 global_mesh.remove_edge(self_node, ne_node)
398 if global_mesh.nodes[self_node].get('far') is None:
399 global_mesh.nodes[self_node]['far'] = []
400 global_mesh.nodes[self_node]['far'].append(ne_node)
401
402 global_fpath_map = np.zeros_like(other_edge_with_id) - 1
403 global_fpath_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] = fpath_map
404 fpath_ids = np.unique(global_fpath_map)
405 fpath_ids = fpath_ids[1:] if fpath_ids.shape[0] > 0 and fpath_ids[0] == -1 else []
406 fpath_real_id_map = np.zeros_like(global_fpath_map) - 1
407 for fpath_id in fpath_ids:
408 fpath_real_id = np.unique(((global_fpath_map == fpath_id).astype(int) * (other_edge_with_id + 1)) - 1)
409 fpath_real_id = fpath_real_id[1:] if fpath_real_id.shape[0] > 0 and fpath_real_id[0] == -1 else []
410 fpath_real_id = fpath_real_id.astype(int)
411 fpath_real_id = np.bincount(fpath_real_id).argmax()
412 fpath_real_id_map[global_fpath_map == fpath_id] = fpath_real_id
413 nxs, nys = np.where((fpath_map > -1))
414 for nx, ny in zip(nxs, nys):
415 self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
416 n_id = fpath_map[nx, ny]
417 four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
418 if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
419 for nex, ney in four_nes:
420 ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
421 if global_mesh.has_node(ne_node) is False:
422 continue
423 if npath_map[nex, ney] == n_id or global_mesh.nodes[ne_node].get('edge_id') == n_id:
424 if global_mesh.has_edge(self_node, ne_node) is True:
425 global_mesh.remove_edge(self_node, ne_node)
426 if global_mesh.nodes[self_node].get('near') is None:
427 global_mesh.nodes[self_node]['near'] = []
428 if global_mesh.nodes[self_node].get('edge_id') is None:
429 f_id = int(round(fpath_real_id_map[self_node[0], self_node[1]]))
430 global_mesh.nodes[self_node]['edge_id'] = f_id
431 info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = f_id
432 edge_ccs[f_id].add(self_node)
433 global_mesh.nodes[self_node]['near'].append(ne_node)
434
435 return info_on_pix, global_mesh, image, depth, edge_ccs
436 # for edge_cc in edge_ccs:
437 # for edge_node in edge_cc:
438 # edge_ccs
439 # context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, init_mask_connect, edge_maps, extend_context_ccs, extend_edge_ccs
440
441def get_valid_size(imap):
442 x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1
443 x_min = np.where(imap.sum(1).squeeze() > 0)[0].min()
444 y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1
445 y_min = np.where(imap.sum(0).squeeze() > 0)[0].min()
446 size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min}
447
448 return size_dict
449
450def dilate_valid_size(isize_dict, imap, dilate=[0, 0]):
451 osize_dict = copy.deepcopy(isize_dict)
452 osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0])
453 osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0])
454 osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0])
455 osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1])
456
457 return osize_dict
458
459def size_operation(size_a, size_b, operation):
460 assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
461 osize = {}
462 if operation == '+':
463 osize['x_min'] = min(size_a['x_min'], size_b['x_min'])
464 osize['y_min'] = min(size_a['y_min'], size_b['y_min'])
465 osize['x_max'] = max(size_a['x_max'], size_b['x_max'])
466 osize['y_max'] = max(size_a['y_max'], size_b['y_max'])
467 assert operation != '-', "Operation '-' is undefined !"
468
469 return osize
470
471def fill_dummy_bord(mesh, info_on_pix, image, depth, config):
472 context = np.zeros_like(depth).astype(np.uint8)
473 context[mesh.graph['hoffset']:mesh.graph['hoffset'] + mesh.graph['noext_H'],
474 mesh.graph['woffset']:mesh.graph['woffset'] + mesh.graph['noext_W']] = 1
475 mask = 1 - context
476 xs, ys = np.where(mask > 0)
477 depth = depth * context
478 image = image * context[..., None]
479 cur_depth = 0
480 cur_disp = 0
481 color = [0, 0, 0]
482 for x, y in zip(xs, ys):
483 cur_node = (x, y, cur_depth)
484 mesh.add_node(cur_node, color=color,
485 synthesis=False,
486 disp=cur_disp,
487 cc_id=set(),
488 ext_pixel=True)
489 info_on_pix[(x, y)] = [{'depth':cur_depth,
490 'color':mesh.nodes[(x, y, cur_depth)]['color'],
491 'synthesis':False,
492 'disp':mesh.nodes[cur_node]['disp'],
493 'ext_pixel':True}]
494 # for x, y in zip(xs, ys):
495 four_nes = [(xx, yy) for xx, yy in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)] if\
496 0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and info_on_pix.get((xx, yy)) is not None]
497 for ne in four_nes:
498 # if (ne[0] - x) + (ne[1] - y) == 1 and info_on_pix.get((ne[0], ne[1])) is not None:
499 mesh.add_edge(cur_node, (ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']))
500
501 return mesh, info_on_pix
502
503
504def enlarge_border(mesh, info_on_pix, depth, image, config):
505 mesh.graph['hoffset'], mesh.graph['woffset'] = config['extrapolation_thickness'], config['extrapolation_thickness']
506 mesh.graph['bord_up'], mesh.graph['bord_left'], mesh.graph['bord_down'], mesh.graph['bord_right'] = \
507 0, 0, mesh.graph['H'], mesh.graph['W']
508 # new_image = np.pad(image,
509 # pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
510 # (config['extrapolation_thickness'], config['extrapolation_thickness']), (0, 0)),
511 # mode='constant')
512 # new_depth = np.pad(depth,
513 # pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
514 # (config['extrapolation_thickness'], config['extrapolation_thickness'])),
515 # mode='constant')
516
517 return mesh, info_on_pix, depth, image
518
519def fill_missing_node(mesh, info_on_pix, image, depth):
520 for x in range(mesh.graph['bord_up'], mesh.graph['bord_down']):
521 for y in range(mesh.graph['bord_left'], mesh.graph['bord_right']):
522 if info_on_pix.get((x, y)) is None:
523 print("fill missing node = ", x, y)
524 #import pdb; pdb.set_trace()
525 re_depth, re_count = 0, 0
526 for ne in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
527 if info_on_pix.get(ne) is not None:
528 re_depth += info_on_pix[ne][0]['depth']
529 re_count += 1
530 if re_count == 0:
531 re_depth = -abs(depth[x, y])
532 else:
533 re_depth = re_depth / re_count
534 depth[x, y] = abs(re_depth)
535 info_on_pix[(x, y)] = [{'depth':re_depth,
536 'color':image[x, y],
537 'synthesis':False,
538 'disp':1./re_depth}]
539 mesh.add_node((x, y, re_depth), color=image[x, y],
540 synthesis=False,
541 disp=1./re_depth,
542 cc_id=set())
543 return mesh, info_on_pix, depth
544
545
546
547def refresh_bord_depth(mesh, info_on_pix, image, depth):
548 H, W = mesh.graph['H'], mesh.graph['W']
549 corner_nodes = [(mesh.graph['bord_up'], mesh.graph['bord_left']),
550 (mesh.graph['bord_up'], mesh.graph['bord_right'] - 1),
551 (mesh.graph['bord_down'] - 1, mesh.graph['bord_left']),
552 (mesh.graph['bord_down'] - 1, mesh.graph['bord_right'] - 1)]
553 # (0, W - 1), (H - 1, 0), (H - 1, W - 1)]
554 bord_nodes = []
555 bord_nodes += [(mesh.graph['bord_up'], xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
556 bord_nodes += [(mesh.graph['bord_down'] - 1, xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
557 bord_nodes += [(xx, mesh.graph['bord_left']) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
558 bord_nodes += [(xx, mesh.graph['bord_right'] - 1) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
559 for xy in bord_nodes:
560 tgt_loc = None
561 if xy[0] == mesh.graph['bord_up']:
562 tgt_loc = (xy[0] + 1, xy[1])# (1, xy[1])
563 elif xy[0] == mesh.graph['bord_down'] - 1:
564 tgt_loc = (xy[0] - 1, xy[1]) # (H - 2, xy[1])
565 elif xy[1] == mesh.graph['bord_left']:
566 tgt_loc = (xy[0], xy[1] + 1)
567 elif xy[1] == mesh.graph['bord_right'] - 1:
568 tgt_loc = (xy[0], xy[1] - 1)
569 if tgt_loc is not None:
570 ne_infos = info_on_pix.get(tgt_loc)
571 if ne_infos is None:
572 import pdb; pdb.set_trace()
573 # if ne_infos is not None and len(ne_infos) == 1:
574 tgt_depth = ne_infos[0]['depth']
575 tgt_disp = ne_infos[0]['disp']
576 new_node = (xy[0], xy[1], tgt_depth)
577 src_node = (tgt_loc[0], tgt_loc[1], tgt_depth)
578 tgt_nes_loc = [(xx[0], xx[1]) \
579 for xx in mesh.neighbors(src_node)]
580 tgt_nes_loc = [(xx[0] - tgt_loc[0] + xy[0], xx[1] - tgt_loc[1] + xy[1]) for xx in tgt_nes_loc \
581 if abs(xx[0] - xy[0]) == 1 and abs(xx[1] - xy[1]) == 1]
582 tgt_nes_loc = [xx for xx in tgt_nes_loc if info_on_pix.get(xx) is not None]
583 tgt_nes_loc.append(tgt_loc)
584 # if (xy[0], xy[1]) == (559, 60):
585 # import pdb; pdb.set_trace()
586 if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
587 old_depth = info_on_pix[xy][0].get('depth')
588 old_node = (xy[0], xy[1], old_depth)
589 mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
590 mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), old_node) for zz in tgt_nes_loc])
591 mapping_dict = {old_node: new_node}
592 # if old_node[2] == new_node[2]:
593 # print("mapping_dict = ", mapping_dict)
594 info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
595 else:
596 info_on_pix[xy] = []
597 info_on_pix[xy][0] = info_on_pix[tgt_loc][0]
598 info_on_pix['color'] = image[xy[0], xy[1]]
599 info_on_pix['old_color'] = image[xy[0], xy[1]]
600 mesh.add_node(new_node)
601 mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), new_node) for zz in tgt_nes_loc])
602 mesh.nodes[new_node]['far'] = None
603 mesh.nodes[new_node]['near'] = None
604 if mesh.nodes[src_node].get('far') is not None:
605 redundant_nodes = [ne for ne in mesh.nodes[src_node]['far'] if (ne[0], ne[1]) == xy]
606 [mesh.nodes[src_node]['far'].remove(aa) for aa in redundant_nodes]
607 if mesh.nodes[src_node].get('near') is not None:
608 redundant_nodes = [ne for ne in mesh.nodes[src_node]['near'] if (ne[0], ne[1]) == xy]
609 [mesh.nodes[src_node]['near'].remove(aa) for aa in redundant_nodes]
610 for xy in corner_nodes:
611 hx, hy = xy
612 four_nes = [xx for xx in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if \
613 mesh.graph['bord_up'] <= xx[0] < mesh.graph['bord_down'] and \
614 mesh.graph['bord_left'] <= xx[1] < mesh.graph['bord_right']]
615 ne_nodes = []
616 ne_depths = []
617 for ne_loc in four_nes:
618 if info_on_pix.get(ne_loc) is not None:
619 ne_depths.append(info_on_pix[ne_loc][0]['depth'])
620 ne_nodes.append((ne_loc[0], ne_loc[1], info_on_pix[ne_loc][0]['depth']))
621 new_node = (xy[0], xy[1], float(np.mean(ne_depths)))
622 if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
623 old_depth = info_on_pix[xy][0].get('depth')
624 old_node = (xy[0], xy[1], old_depth)
625 mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
626 mesh.add_edges_from([(zz, old_node) for zz in ne_nodes])
627 mapping_dict = {old_node: new_node}
628 info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
629 else:
630 info_on_pix[xy] = []
631 info_on_pix[xy][0] = info_on_pix[ne_loc[-1]][0]
632 info_on_pix['color'] = image[xy[0], xy[1]]
633 info_on_pix['old_color'] = image[xy[0], xy[1]]
634 mesh.add_node(new_node)
635 mesh.add_edges_from([(zz, new_node) for zz in ne_nodes])
636 mesh.nodes[new_node]['far'] = None
637 mesh.nodes[new_node]['near'] = None
638 for xy in bord_nodes + corner_nodes:
639 # if (xy[0], xy[1]) == (559, 60):
640 # import pdb; pdb.set_trace()
641 depth[xy[0], xy[1]] = abs(info_on_pix[xy][0]['depth'])
642 for xy in bord_nodes:
643 cur_node = (xy[0], xy[1], info_on_pix[xy][0]['depth'])
644 nes = mesh.neighbors(cur_node)
645 four_nes = set([(xy[0] + 1, xy[1]), (xy[0] - 1, xy[1]), (xy[0], xy[1] + 1), (xy[0], xy[1] - 1)]) - \
646 set([(ne[0], ne[1]) for ne in nes])
647 four_nes = [ne for ne in four_nes if mesh.graph['bord_up'] <= ne[0] < mesh.graph['bord_down'] and \
648 mesh.graph['bord_left'] <= ne[1] < mesh.graph['bord_right']]
649 four_nes = [(ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']) for ne in four_nes]
650 mesh.nodes[cur_node]['far'] = []
651 mesh.nodes[cur_node]['near'] = []
652 for ne in four_nes:
653 if abs(ne[2]) >= abs(cur_node[2]):
654 mesh.nodes[cur_node]['far'].append(ne)
655 else:
656 mesh.nodes[cur_node]['near'].append(ne)
657
658 return mesh, info_on_pix, depth
659
660def get_union_size(mesh, dilate, *alls_cc):
661 all_cc = reduce(lambda x, y: x | y, [set()] + [*alls_cc])
662 min_x, min_y, max_x, max_y = mesh.graph['H'], mesh.graph['W'], 0, 0
663 H, W = mesh.graph['H'], mesh.graph['W']
664 for node in all_cc:
665 if node[0] < min_x:
666 min_x = node[0]
667 if node[0] > max_x:
668 max_x = node[0]
669 if node[1] < min_y:
670 min_y = node[1]
671 if node[1] > max_y:
672 max_y = node[1]
673 max_x = max_x + 1
674 max_y = max_y + 1
675 # mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20])
676 osize_dict = dict()
677 osize_dict['x_min'] = max(0, min_x - dilate[0])
678 osize_dict['x_max'] = min(H, max_x + dilate[0])
679 osize_dict['y_min'] = max(0, min_y - dilate[1])
680 osize_dict['y_max'] = min(W, max_y + dilate[1])
681
682 return osize_dict
683
684def incomplete_node(mesh, edge_maps, info_on_pix):
685 vis_map = np.zeros((mesh.graph['H'], mesh.graph['W']))
686
687 for node in mesh.nodes:
688 if mesh.nodes[node].get('synthesis') is not True:
689 connect_all_flag = False
690 nes = [xx for xx in mesh.neighbors(node) if mesh.nodes[xx].get('synthesis') is not True]
691 if len(nes) < 3 and 0 < node[0] < mesh.graph['H'] - 1 and 0 < node[1] < mesh.graph['W'] - 1:
692 if len(nes) <= 1:
693 connect_all_flag = True
694 else:
695 dan_ne_node_a = nes[0]
696 dan_ne_node_b = nes[1]
697 if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \
698 abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1:
699 connect_all_flag = True
700 if connect_all_flag == True:
701 vis_map[node[0], node[1]] = len(nes)
702 four_nes = [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)]
703 for ne in four_nes:
704 for info in info_on_pix[(ne[0], ne[1])]:
705 ne_node = (ne[0], ne[1], info['depth'])
706 if info.get('synthesis') is not True and mesh.has_node(ne_node):
707 mesh.add_edge(node, ne_node)
708 break
709
710 return mesh
711
712def edge_inpainting(edge_id, context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
713 mesh, edge_map, edge_maps_with_id, config, union_size, depth_edge_model, inpaint_iter):
714 edge_dict = get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
715 mesh.graph['H'], mesh.graph['W'], mesh)
716 edge_dict['edge'], end_depth_maps, _ = \
717 filter_irrelevant_edge_new(edge_dict['self_edge'] + edge_dict['comp_edge'],
718 edge_map,
719 edge_maps_with_id,
720 edge_id,
721 edge_dict['context'],
722 edge_dict['depth'], mesh, context_cc | erode_context_cc, spdb=True)
723 patch_edge_dict = dict()
724 patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \
725 patch_edge_dict['disp'], patch_edge_dict['edge'] = \
726 crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'],
727 edge_dict['rgb'], edge_dict['disp'], edge_dict['edge'])
728 tensor_edge_dict = convert2tensor(patch_edge_dict)
729 if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
730 with torch.no_grad():
731 device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
732 depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'],
733 tensor_edge_dict['context'],
734 tensor_edge_dict['rgb'],
735 tensor_edge_dict['disp'],
736 tensor_edge_dict['edge'],
737 unit_length=128,
738 cuda=device)
739 depth_edge_output = depth_edge_output.cpu()
740 tensor_edge_dict['output'] = (depth_edge_output > config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge']
741 else:
742 tensor_edge_dict['output'] = tensor_edge_dict['edge']
743 depth_edge_output = tensor_edge_dict['edge'] + 0
744 patch_edge_dict['output'] = tensor_edge_dict['output'].squeeze().data.cpu().numpy()
745 edge_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
746 edge_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
747 patch_edge_dict['output']
748
749 return edge_dict, end_depth_maps
750
751def depth_inpainting(context_cc, extend_context_cc, erode_context_cc, mask_cc, mesh, config, union_size, depth_feat_model, edge_output, given_depth_dict=False, spdb=False):
752 if given_depth_dict is False:
753 depth_dict = get_depth_from_nodes(context_cc | extend_context_cc, erode_context_cc, mask_cc, mesh.graph['H'], mesh.graph['W'], mesh, config['log_depth'])
754 if edge_output is not None:
755 depth_dict['edge'] = edge_output
756 else:
757 depth_dict = given_depth_dict
758 patch_depth_dict = dict()
759 patch_depth_dict['mask'], patch_depth_dict['context'], patch_depth_dict['depth'], \
760 patch_depth_dict['zero_mean_depth'], patch_depth_dict['edge'] = \
761 crop_maps_by_size(union_size, depth_dict['mask'], depth_dict['context'],
762 depth_dict['real_depth'], depth_dict['zero_mean_depth'], depth_dict['edge'])
763 tensor_depth_dict = convert2tensor(patch_depth_dict)
764 resize_mask = open_small_mask(tensor_depth_dict['mask'], tensor_depth_dict['context'], 3, 41)
765 with torch.no_grad():
766 device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
767 depth_output = depth_feat_model.forward_3P(resize_mask,
768 tensor_depth_dict['context'],
769 tensor_depth_dict['zero_mean_depth'],
770 tensor_depth_dict['edge'],
771 unit_length=128,
772 cuda=device)
773 depth_output = depth_output.cpu()
774 tensor_depth_dict['output'] = torch.exp(depth_output + depth_dict['mean_depth']) * \
775 tensor_depth_dict['mask'] + tensor_depth_dict['depth']
776 patch_depth_dict['output'] = tensor_depth_dict['output'].data.cpu().numpy().squeeze()
777 depth_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
778 depth_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
779 patch_depth_dict['output']
780 depth_output = depth_dict['output'] * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']
781 depth_output = smooth_cntsyn_gap(depth_dict['output'].copy() * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context'],
782 depth_dict['mask'], depth_dict['context'],
783 init_mask_region=depth_dict['mask'])
784 if spdb is True:
785 f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True);
786 ax1.imshow(depth_output * depth_dict['mask'] + depth_dict['depth']); ax2.imshow(depth_dict['output'] * depth_dict['mask'] + depth_dict['depth']); plt.show()
787 import pdb; pdb.set_trace()
788 depth_dict['output'] = depth_output * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']
789
790 return depth_dict
791
792def update_info(mapping_dict, info_on_pix, *meshes):
793 rt_meshes = []
794 for mesh in meshes:
795 rt_meshes.append(relabel_node(mesh, mesh.nodes, [*mapping_dict.keys()][0], [*mapping_dict.values()][0]))
796 x, y, _ = [*mapping_dict.keys()][0]
797 info_on_pix[(x, y)][0]['depth'] = [*mapping_dict.values()][0][2]
798
799 return [info_on_pix] + rt_meshes
800
801def build_connection(mesh, cur_node, dst_node):
802 if (abs(cur_node[0] - dst_node[0]) + abs(cur_node[1] - dst_node[1])) < 2:
803 mesh.add_edge(cur_node, dst_node)
804 if abs(cur_node[0] - dst_node[0]) > 1 or abs(cur_node[1] - dst_node[1]) > 1:
805 return mesh
806 ne_nodes = [*mesh.neighbors(cur_node)].copy()
807 for ne_node in ne_nodes:
808 if mesh.has_edge(ne_node, dst_node) or ne_node == dst_node:
809 continue
810 else:
811 mesh = build_connection(mesh, ne_node, dst_node)
812
813 return mesh
814
815def recursive_add_edge(edge_mesh, mesh, info_on_pix, cur_node, mark):
816 ne_nodes = [(x[0], x[1]) for x in edge_mesh.neighbors(cur_node)]
817 for node_xy in ne_nodes:
818 node = (node_xy[0], node_xy[1], info_on_pix[node_xy][0]['depth'])
819 if mark[node[0], node[1]] != 3:
820 continue
821 else:
822 mark[node[0], node[1]] = 0
823 mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)])
824 mesh = build_connection(mesh, cur_node, node)
825 re_info = dict(depth=0, count=0)
826 for re_ne in mesh.neighbors(node):
827 re_info['depth'] += re_ne[2]
828 re_info['count'] += 1.
829 try:
830 re_depth = re_info['depth'] / re_info['count']
831 except:
832 re_depth = node[2]
833 re_node = (node_xy[0], node_xy[1], re_depth)
834 mapping_dict = {node: re_node}
835 info_on_pix, edge_mesh, mesh = update_info(mapping_dict, info_on_pix, edge_mesh, mesh)
836
837 edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, re_node, mark)
838
839 return edge_mesh, mesh, mark, info_on_pix
840
841def resize_for_edge(tensor_dict, largest_size):
842 resize_dict = {k: v.clone() for k, v in tensor_dict.items()}
843 frac = largest_size / np.array([*resize_dict['edge'].shape[-2:]]).max()
844 if frac < 1:
845 resize_mark = torch.nn.functional.interpolate(torch.cat((resize_dict['mask'],
846 resize_dict['context']),
847 dim=1),
848 scale_factor=frac,
849 mode='bilinear')
850 resize_dict['mask'] = (resize_mark[:, 0:1] > 0).float()
851 resize_dict['context'] = (resize_mark[:, 1:2] == 1).float()
852 resize_dict['context'][resize_dict['mask'] > 0] = 0
853 resize_dict['edge'] = torch.nn.functional.interpolate(resize_dict['edge'],
854 scale_factor=frac,
855 mode='bilinear')
856 resize_dict['edge'] = (resize_dict['edge'] > 0).float()
857 resize_dict['edge'] = resize_dict['edge'] * resize_dict['context']
858 resize_dict['disp'] = torch.nn.functional.interpolate(resize_dict['disp'],
859 scale_factor=frac,
860 mode='nearest')
861 resize_dict['disp'] = resize_dict['disp'] * resize_dict['context']
862 resize_dict['rgb'] = torch.nn.functional.interpolate(resize_dict['rgb'],
863 scale_factor=frac,
864 mode='bilinear')
865 resize_dict['rgb'] = resize_dict['rgb'] * resize_dict['context']
866 return resize_dict
867
868def get_map_from_nodes(nodes, height, width):
869 omap = np.zeros((height, width))
870 for n in nodes:
871 omap[n[0], n[1]] = 1
872
873 return omap
874
875def get_map_from_ccs(ccs, height, width, condition_input=None, condition=None, real_id=False, id_shift=0):
876 if condition is None:
877 condition = lambda x, condition_input: True
878
879 if real_id is True:
880 omap = np.zeros((height, width)) + (-1) + id_shift
881 else:
882 omap = np.zeros((height, width))
883 for cc_id, cc in enumerate(ccs):
884 for n in cc:
885 if condition(n, condition_input):
886 if real_id is True:
887 omap[n[0], n[1]] = cc_id + id_shift
888 else:
889 omap[n[0], n[1]] = 1
890 return omap
891
892def revise_map_by_nodes(nodes, imap, operation, limit_constr=None):
893 assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
894 omap = copy.deepcopy(imap)
895 revise_flag = True
896 if operation == '+':
897 for n in nodes:
898 omap[n[0], n[1]] = 1
899 if limit_constr is not None and omap.sum() > limit_constr:
900 omap = imap
901 revise_flag = False
902 elif operation == '-':
903 for n in nodes:
904 omap[n[0], n[1]] = 0
905 if limit_constr is not None and omap.sum() < limit_constr:
906 omap = imap
907 revise_flag = False
908
909 return omap, revise_flag
910
911def repaint_info(mesh, cc, x_anchor, y_anchor, source_type):
912 if source_type == 'rgb':
913 feat = np.zeros((3, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
914 else:
915 feat = np.zeros((1, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
916 for node in cc:
917 if source_type == 'rgb':
918 feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = np.array(mesh.nodes[node]['color']) / 255.
919 elif source_type == 'd':
920 feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = abs(node[2])
921
922 return feat
923
924def get_context_from_nodes(mesh, cc, H, W, source_type=''):
925 if 'rgb' in source_type or 'color' in source_type:
926 feat = np.zeros((H, W, 3))
927 else:
928 feat = np.zeros((H, W))
929 context = np.zeros((H, W))
930 for node in cc:
931 if 'rgb' in source_type or 'color' in source_type:
932 feat[node[0], node[1]] = np.array(mesh.nodes[node]['color']) / 255.
933 context[node[0], node[1]] = 1
934 else:
935 feat[node[0], node[1]] = abs(node[2])
936
937 return feat, context
938
939def get_mask_from_nodes(mesh, cc, H, W):
940 mask = np.zeros((H, W))
941 for node in cc:
942 mask[node[0], node[1]] = abs(node[2])
943
944 return mask
945
946
947def get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc, H, W, mesh):
948 context = np.zeros((H, W))
949 mask = np.zeros((H, W))
950 rgb = np.zeros((H, W, 3))
951 disp = np.zeros((H, W))
952 depth = np.zeros((H, W))
953 real_depth = np.zeros((H, W))
954 edge = np.zeros((H, W))
955 comp_edge = np.zeros((H, W))
956 fpath_map = np.zeros((H, W)) - 1
957 npath_map = np.zeros((H, W)) - 1
958 near_depth = np.zeros((H, W))
959 for node in context_cc:
960 rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
961 disp[node[0], node[1]] = mesh.nodes[node]['disp']
962 depth[node[0], node[1]] = node[2]
963 context[node[0], node[1]] = 1
964 for node in erode_context_cc:
965 rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
966 disp[node[0], node[1]] = mesh.nodes[node]['disp']
967 depth[node[0], node[1]] = node[2]
968 context[node[0], node[1]] = 1
969 rgb = rgb / 255.
970 disp = np.abs(disp)
971 disp = disp / disp.max()
972 real_depth = depth.copy()
973 for node in context_cc:
974 if mesh.nodes[node].get('real_depth') is not None:
975 real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
976 for node in erode_context_cc:
977 if mesh.nodes[node].get('real_depth') is not None:
978 real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
979 for node in mask_cc:
980 mask[node[0], node[1]] = 1
981 near_depth[node[0], node[1]] = node[2]
982 for node in edge_cc:
983 edge[node[0], node[1]] = 1
984 for node in extend_edge_cc:
985 comp_edge[node[0], node[1]] = 1
986 rt_dict = {'rgb': rgb, 'disp': disp, 'depth': depth, 'real_depth': real_depth, 'self_edge': edge, 'context': context,
987 'mask': mask, 'fpath_map': fpath_map, 'npath_map': npath_map, 'comp_edge': comp_edge, 'valid_area': context + mask,
988 'near_depth': near_depth}
989
990 return rt_dict
991
992def get_depth_from_maps(context_map, mask_map, depth_map, H, W, log_depth=False):
993 context = context_map.astype(np.uint8)
994 mask = mask_map.astype(np.uint8).copy()
995 depth = np.abs(depth_map)
996 real_depth = depth.copy()
997 zero_mean_depth = np.zeros((H, W))
998
999 if log_depth is True:
1000 log_depth = np.log(real_depth + 1e-8) * context
1001 mean_depth = np.mean(log_depth[context > 0])
1002 zero_mean_depth = (log_depth - mean_depth) * context
1003 else:
1004 zero_mean_depth = real_depth
1005 mean_depth = 0
1006 edge = np.zeros_like(depth)
1007
1008 rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
1009 'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth, 'edge': edge}
1010
1011 return rt_dict
1012
1013def get_depth_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh, log_depth=False):
1014 context = np.zeros((H, W))
1015 mask = np.zeros((H, W))
1016 depth = np.zeros((H, W))
1017 real_depth = np.zeros((H, W))
1018 zero_mean_depth = np.zeros((H, W))
1019 for node in context_cc:
1020 depth[node[0], node[1]] = node[2]
1021 context[node[0], node[1]] = 1
1022 for node in erode_context_cc:
1023 depth[node[0], node[1]] = node[2]
1024 context[node[0], node[1]] = 1
1025 depth = np.abs(depth)
1026 real_depth = depth.copy()
1027 for node in context_cc:
1028 if mesh.nodes[node].get('real_depth') is not None:
1029 real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
1030 for node in erode_context_cc:
1031 if mesh.nodes[node].get('real_depth') is not None:
1032 real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
1033 real_depth = np.abs(real_depth)
1034 for node in mask_cc:
1035 mask[node[0], node[1]] = 1
1036 if log_depth is True:
1037 log_depth = np.log(real_depth + 1e-8) * context
1038 mean_depth = np.mean(log_depth[context > 0])
1039 zero_mean_depth = (log_depth - mean_depth) * context
1040 else:
1041 zero_mean_depth = real_depth
1042 mean_depth = 0
1043
1044 rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
1045 'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth}
1046
1047 return rt_dict
1048
1049def get_rgb_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh):
1050 context = np.zeros((H, W))
1051 mask = np.zeros((H, W))
1052 rgb = np.zeros((H, W, 3))
1053 erode_context = np.zeros((H, W))
1054 for node in context_cc:
1055 rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
1056 context[node[0], node[1]] = 1
1057 rgb = rgb / 255.
1058 for node in mask_cc:
1059 mask[node[0], node[1]] = 1
1060 for node in erode_context_cc:
1061 erode_context[node[0], node[1]] = 1
1062 mask[node[0], node[1]] = 1
1063 rt_dict = {'rgb': rgb, 'context': context, 'mask': mask,
1064 'erode': erode_context}
1065
1066 return rt_dict
1067
1068def crop_maps_by_size(size, *imaps):
1069 omaps = []
1070 for imap in imaps:
1071 omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy())
1072
1073 return omaps
1074
1075def convert2tensor(input_dict):
1076 rt_dict = {}
1077 for key, value in input_dict.items():
1078 if 'rgb' in key or 'color' in key:
1079 rt_dict[key] = torch.FloatTensor(value).permute(2, 0, 1)[None, ...]
1080 else:
1081 rt_dict[key] = torch.FloatTensor(value)[None, None, ...]
1082
1083 return rt_dict
1084