Spaces:
Runtime error
Runtime error
| import os | |
| import glob | |
| import cv2 | |
| import scipy.misc as misc | |
| from skimage.transform import resize | |
| import numpy as np | |
| from functools import reduce | |
| from operator import mul | |
| import torch | |
| from torch import nn | |
| import matplotlib.pyplot as plt | |
| import re | |
| try: | |
| import cynetworkx as netx | |
| except ImportError: | |
| import networkx as netx | |
| from scipy.ndimage import gaussian_filter | |
| from skimage.feature import canny | |
| import collections | |
| import shutil | |
| import imageio | |
| import copy | |
| from matplotlib import pyplot as plt | |
| from mpl_toolkits.mplot3d import Axes3D | |
| import time | |
| from scipy.interpolate import interp1d | |
| from collections import namedtuple | |
| def path_planning(num_frames, x, y, z, path_type=''): | |
| if path_type == 'straight-line': | |
| corner_points = np.array([[0, 0, 0], [(0 + x) * 0.5, (0 + y) * 0.5, (0 + z) * 0.5], [x, y, z]]) | |
| corner_t = np.linspace(0, 1, len(corner_points)) | |
| t = np.linspace(0, 1, num_frames) | |
| cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic') | |
| spline = cs(t) | |
| xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)] | |
| elif path_type == 'double-straight-line': | |
| corner_points = np.array([[-x, -y, -z], [0, 0, 0], [x, y, z]]) | |
| corner_t = np.linspace(0, 1, len(corner_points)) | |
| t = np.linspace(0, 1, num_frames) | |
| cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic') | |
| spline = cs(t) | |
| xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)] | |
| elif path_type == 'circle': | |
| xs, ys, zs = [], [], [] | |
| for frame_id, bs_shift_val in enumerate(np.arange(-2.0, 2.0, (4./num_frames))): | |
| xs += [np.cos(bs_shift_val * np.pi) * 1 * x] | |
| ys += [np.sin(bs_shift_val * np.pi) * 1 * y] | |
| zs += [np.cos(bs_shift_val * np.pi/2.) * 1 * z] | |
| xs, ys, zs = np.array(xs), np.array(ys), np.array(zs) | |
| return xs, ys, zs | |
| def open_small_mask(mask, context, open_iteration, kernel): | |
| np_mask = mask.cpu().data.numpy().squeeze().astype(np.uint8) | |
| raw_mask = np_mask.copy() | |
| np_context = context.cpu().data.numpy().squeeze().astype(np.uint8) | |
| np_input = np_mask + np_context | |
| for _ in range(open_iteration): | |
| np_input = cv2.erode(cv2.dilate(np_input, np.ones((kernel, kernel)), iterations=1), np.ones((kernel,kernel)), iterations=1) | |
| np_mask[(np_input - np_context) > 0] = 1 | |
| out_mask = torch.FloatTensor(np_mask).to(mask)[None, None, ...] | |
| return out_mask | |
| def filter_irrelevant_edge_new(self_edge, comp_edge, other_edges, other_edges_with_id, current_edge_id, context, depth, mesh, context_cc, spdb=False): | |
| other_edges = other_edges.squeeze().astype(np.uint8) | |
| other_edges_with_id = other_edges_with_id.squeeze() | |
| self_edge = self_edge.squeeze() | |
| dilate_bevel_self_edge = cv2.dilate((self_edge + comp_edge).astype(np.uint8), np.array([[1,1,1],[1,1,1],[1,1,1]]), iterations=1) | |
| dilate_cross_self_edge = cv2.dilate((self_edge + comp_edge).astype(np.uint8), np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1) | |
| edge_ids = np.unique(other_edges_with_id * context + (-1) * (1 - context)).astype(np.int) | |
| end_depth_maps = np.zeros_like(self_edge) | |
| self_edge_ids = np.sort(np.unique(other_edges_with_id[self_edge > 0]).astype(np.int)) | |
| self_edge_ids = self_edge_ids[1:] if self_edge_ids.shape[0] > 0 and self_edge_ids[0] == -1 else self_edge_ids | |
| self_comp_ids = np.sort(np.unique(other_edges_with_id[comp_edge > 0]).astype(np.int)) | |
| self_comp_ids = self_comp_ids[1:] if self_comp_ids.shape[0] > 0 and self_comp_ids[0] == -1 else self_comp_ids | |
| edge_ids = edge_ids[1:] if edge_ids[0] == -1 else edge_ids | |
| other_edges_info = [] | |
| extend_other_edges = np.zeros_like(other_edges) | |
| if spdb is True: | |
| f, ((ax1, ax2, ax3)) = plt.subplots(1, 3, sharex=True, sharey=True); ax1.imshow(self_edge); ax2.imshow(context); ax3.imshow(other_edges_with_id * context + (-1) * (1 - context)); plt.show() | |
| import pdb; pdb.set_trace() | |
| filter_self_edge = np.zeros_like(self_edge) | |
| for self_edge_id in self_edge_ids: | |
| filter_self_edge[other_edges_with_id == self_edge_id] = 1 | |
| dilate_self_comp_edge = cv2.dilate(comp_edge, kernel=np.ones((3, 3)), iterations=2) | |
| valid_self_comp_edge = np.zeros_like(comp_edge) | |
| for self_comp_id in self_comp_ids: | |
| valid_self_comp_edge[self_comp_id == other_edges_with_id] = 1 | |
| self_comp_edge = dilate_self_comp_edge * valid_self_comp_edge | |
| filter_self_edge = (filter_self_edge + self_comp_edge).clip(0, 1) | |
| for edge_id in edge_ids: | |
| other_edge_locs = (other_edges_with_id == edge_id).astype(np.uint8) | |
| condition = (other_edge_locs * other_edges * context.astype(np.uint8)) | |
| end_cross_point = dilate_cross_self_edge * condition * (1 - filter_self_edge) | |
| end_bevel_point = dilate_bevel_self_edge * condition * (1 - filter_self_edge) | |
| if end_bevel_point.max() != 0: | |
| end_depth_maps[end_bevel_point != 0] = depth[end_bevel_point != 0] | |
| if end_cross_point.max() == 0: | |
| nxs, nys = np.where(end_bevel_point != 0) | |
| for nx, ny in zip(nxs, nys): | |
| bevel_node = [xx for xx in context_cc if xx[0] == nx and xx[1] == ny][0] | |
| for ne in mesh.neighbors(bevel_node): | |
| if other_edges_with_id[ne[0], ne[1]] > -1 and dilate_cross_self_edge[ne[0], ne[1]] > 0: | |
| extend_other_edges[ne[0], ne[1]] = 1 | |
| break | |
| else: | |
| other_edges[other_edges_with_id == edge_id] = 0 | |
| other_edges = (other_edges + extend_other_edges).clip(0, 1) * context | |
| return other_edges, end_depth_maps, other_edges_info | |
| def clean_far_edge_new(input_edge, end_depth_maps, mask, context, global_mesh, info_on_pix, self_edge, inpaint_id, config): | |
| mesh = netx.Graph() | |
| hxs, hys = np.where(input_edge * mask > 0) | |
| valid_near_edge = (input_edge != 0).astype(np.uint8) * context | |
| valid_map = mask + context | |
| invalid_edge_ids = [] | |
| for hx, hy in zip(hxs, hys): | |
| node = (hx ,hy) | |
| mesh.add_node((hx, hy)) | |
| eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \ | |
| (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\ | |
| if 0 <= ne[0] < input_edge.shape[0] and 0 <= ne[1] < input_edge.shape[1] and 0 < input_edge[ne[0], ne[1]]] # or end_depth_maps[ne[0], ne[1]] != 0] | |
| for ne in eight_nes: | |
| mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy)) | |
| if end_depth_maps[ne[0], ne[1]] != 0: | |
| mesh.nodes[ne[0], ne[1]]['cnt'] = True | |
| if end_depth_maps[ne[0], ne[1]] == 0: | |
| import pdb; pdb.set_trace() | |
| mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]] | |
| elif mask[ne[0], ne[1]] != 1: | |
| four_nes = [nne for nne in [(ne[0] + 1, ne[1]), (ne[0] - 1, ne[1]), (ne[0], ne[1] + 1), (ne[0], ne[1] - 1)]\ | |
| if nne[0] < end_depth_maps.shape[0] and nne[0] >= 0 and nne[1] < end_depth_maps.shape[1] and nne[1] >= 0] | |
| for nne in four_nes: | |
| if end_depth_maps[nne[0], nne[1]] != 0: | |
| mesh.add_edge(nne, ne, length=np.hypot(nne[0] - ne[0], nne[1] - ne[1])) | |
| mesh.nodes[nne[0], nne[1]]['cnt'] = True | |
| mesh.nodes[nne[0], nne[1]]['depth'] = end_depth_maps[nne[0], nne[1]] | |
| ccs = [*netx.connected_components(mesh)] | |
| end_pts = [] | |
| for cc in ccs: | |
| end_pts.append(set()) | |
| for node in cc: | |
| if mesh.nodes[node].get('cnt') is not None: | |
| end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth'])) | |
| predef_npaths = [None for _ in range(len(ccs))] | |
| fpath_map = np.zeros_like(input_edge) - 1 | |
| npath_map = np.zeros_like(input_edge) - 1 | |
| npaths, fpaths = dict(), dict() | |
| break_flag = False | |
| end_idx = 0 | |
| while end_idx < len(end_pts): | |
| end_pt, cc = [*zip(end_pts, ccs)][end_idx] | |
| end_idx += 1 | |
| sorted_end_pt = [] | |
| fpath = [] | |
| iter_fpath = [] | |
| if len(end_pt) > 2 or len(end_pt) == 0: | |
| if len(end_pt) > 2: | |
| continue | |
| continue | |
| if len(end_pt) == 2: | |
| ravel_end = [*end_pt] | |
| tmp_sub_mesh = mesh.subgraph(list(cc)).copy() | |
| tmp_npath = [*netx.shortest_path(tmp_sub_mesh, (ravel_end[0][0], ravel_end[0][1]), (ravel_end[1][0], ravel_end[1][1]), weight='length')] | |
| fpath_map1, npath_map1, disp_diff1 = plan_path(mesh, info_on_pix, cc, ravel_end[0:1], global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=tmp_npath) | |
| fpath_map2, npath_map2, disp_diff2 = plan_path(mesh, info_on_pix, cc, ravel_end[1:2], global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=tmp_npath) | |
| tmp_disp_diff = [disp_diff1, disp_diff2] | |
| self_end = [] | |
| edge_len = [] | |
| ds_edge = cv2.dilate(self_edge.astype(np.uint8), np.ones((3, 3)), iterations=1) | |
| if ds_edge[ravel_end[0][0], ravel_end[0][1]] > 0: | |
| self_end.append(1) | |
| else: | |
| self_end.append(0) | |
| if ds_edge[ravel_end[1][0], ravel_end[1][1]] > 0: | |
| self_end.append(1) | |
| else: | |
| self_end.append(0) | |
| edge_len = [np.count_nonzero(npath_map1), np.count_nonzero(npath_map2)] | |
| sorted_end_pts = [xx[0] for xx in sorted(zip(ravel_end, self_end, edge_len, [disp_diff1, disp_diff2]), key=lambda x: (x[1], x[2]), reverse=True)] | |
| re_npath_map1, re_fpath_map1 = (npath_map1 != -1).astype(np.uint8), (fpath_map1 != -1).astype(np.uint8) | |
| re_npath_map2, re_fpath_map2 = (npath_map2 != -1).astype(np.uint8), (fpath_map2 != -1).astype(np.uint8) | |
| if np.count_nonzero(re_npath_map1 * re_npath_map2 * mask) / \ | |
| (np.count_nonzero((re_npath_map1 + re_npath_map2) * mask) + 1e-6) > 0.5\ | |
| and np.count_nonzero(re_fpath_map1 * re_fpath_map2 * mask) / \ | |
| (np.count_nonzero((re_fpath_map1 + re_fpath_map2) * mask) + 1e-6) > 0.5\ | |
| and tmp_disp_diff[0] != -1 and tmp_disp_diff[1] != -1: | |
| my_fpath_map, my_npath_map, npath, fpath = \ | |
| plan_path_e2e(mesh, cc, sorted_end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None) | |
| npath_map[my_npath_map != -1] = my_npath_map[my_npath_map != -1] | |
| fpath_map[my_fpath_map != -1] = my_fpath_map[my_fpath_map != -1] | |
| if len(fpath) > 0: | |
| edge_id = global_mesh.nodes[[*sorted_end_pts][0]]['edge_id'] | |
| fpaths[edge_id] = fpath | |
| npaths[edge_id] = npath | |
| invalid_edge_ids.append(edge_id) | |
| else: | |
| if tmp_disp_diff[0] != -1: | |
| ratio_a = tmp_disp_diff[0] / (np.sum(tmp_disp_diff) + 1e-8) | |
| else: | |
| ratio_a = 0 | |
| if tmp_disp_diff[1] != -1: | |
| ratio_b = tmp_disp_diff[1] / (np.sum(tmp_disp_diff) + 1e-8) | |
| else: | |
| ratio_b = 0 | |
| npath_len = len(tmp_npath) | |
| if npath_len > config['depth_edge_dilate_2'] * 2: | |
| npath_len = npath_len - (config['depth_edge_dilate_2'] * 1) | |
| tmp_npath_a = tmp_npath[:int(np.floor(npath_len * ratio_a))] | |
| tmp_npath_b = tmp_npath[::-1][:int(np.floor(npath_len * ratio_b))] | |
| tmp_merge = [] | |
| if len(tmp_npath_a) > 0 and sorted_end_pts[0][0] == tmp_npath_a[0][0] and sorted_end_pts[0][1] == tmp_npath_a[0][1]: | |
| if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0: | |
| tmp_merge.append([sorted_end_pts[:1], tmp_npath_a]) | |
| if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0: | |
| tmp_merge.append([sorted_end_pts[1:2], tmp_npath_b]) | |
| elif len(tmp_npath_b) > 0 and sorted_end_pts[0][0] == tmp_npath_b[0][0] and sorted_end_pts[0][1] == tmp_npath_b[0][1]: | |
| if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0: | |
| tmp_merge.append([sorted_end_pts[:1], tmp_npath_b]) | |
| if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0: | |
| tmp_merge.append([sorted_end_pts[1:2], tmp_npath_a]) | |
| for tmp_idx in range(len(tmp_merge)): | |
| if len(tmp_merge[tmp_idx][1]) == 0: | |
| continue | |
| end_pts.append(tmp_merge[tmp_idx][0]) | |
| ccs.append(set(tmp_merge[tmp_idx][1])) | |
| if len(end_pt) == 1: | |
| sub_mesh = mesh.subgraph(list(cc)).copy() | |
| pnodes = netx.periphery(sub_mesh) | |
| if len(end_pt) == 1: | |
| ends = [*end_pt] | |
| elif len(sorted_end_pt) == 1: | |
| ends = [*sorted_end_pt] | |
| else: | |
| import pdb; pdb.set_trace() | |
| try: | |
| edge_id = global_mesh.nodes[ends[0]]['edge_id'] | |
| except: | |
| import pdb; pdb.set_trace() | |
| pnodes = sorted(pnodes, | |
| key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])), | |
| reverse=True)[0] | |
| npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')] | |
| for np_node in npath: | |
| npath_map[np_node[0], np_node[1]] = edge_id | |
| fpath = [] | |
| if global_mesh.nodes[ends[0]].get('far') is None: | |
| print("None far") | |
| else: | |
| fnodes = global_mesh.nodes[ends[0]].get('far') | |
| dmask = mask + 0 | |
| did = 0 | |
| while True: | |
| did += 1 | |
| dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1) | |
| if did > 3: | |
| break | |
| ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\ | |
| global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)] | |
| if len(ffnode) > 0: | |
| fnode = ffnode[0] | |
| break | |
| if len(ffnode) == 0: | |
| continue | |
| fpath.append((fnode[0], fnode[1])) | |
| barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]]) | |
| n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1])) | |
| while True: | |
| if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]: | |
| n2f_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| for step in range(0, len(npath)): | |
| if step == 0: | |
| continue | |
| elif step == 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]: | |
| next_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| elif step > 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]: | |
| next_barrel = barrel_pair.copy() | |
| break | |
| barrel_pair = np.roll(barrel_pair, 1, axis=1) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| new_locs = [] | |
| if abs(n2f_dir[0]) == 1: | |
| new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1])) | |
| if abs(n2f_dir[1]) == 1: | |
| new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1])) | |
| if len(new_locs) > 1: | |
| new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1]))) | |
| break_flag = False | |
| for new_loc in new_locs: | |
| new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]), | |
| (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\ | |
| if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]] | |
| if np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True: | |
| break | |
| if npath_map[new_loc[0], new_loc[1]] != -1: | |
| if npath_map[new_loc[0], new_loc[1]] != edge_id: | |
| break_flag = True | |
| break | |
| else: | |
| continue | |
| if valid_map[new_loc[0], new_loc[1]] == 0: | |
| break_flag = True | |
| break | |
| fpath.append(new_loc) | |
| if break_flag is True: | |
| break | |
| if step != len(npath) - 1: | |
| for xx in npath[step:]: | |
| if npath_map[xx[0], xx[1]] == edge_id: | |
| npath_map[xx[0], xx[1]] = -1 | |
| npath = npath[:step] | |
| if len(fpath) > 0: | |
| for fp_node in fpath: | |
| fpath_map[fp_node[0], fp_node[1]] = edge_id | |
| fpaths[edge_id] = fpath | |
| npaths[edge_id] = npath | |
| fpath_map[valid_near_edge != 0] = -1 | |
| if len(fpath) > 0: | |
| iter_fpath = copy.deepcopy(fpaths[edge_id]) | |
| for node in iter_fpath: | |
| if valid_near_edge[node[0], node[1]] != 0: | |
| fpaths[edge_id].remove(node) | |
| return fpath_map, npath_map, False, npaths, fpaths, invalid_edge_ids | |
| def plan_path_e2e(mesh, cc, end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None): | |
| my_npath_map = np.zeros_like(input_edge) - 1 | |
| my_fpath_map = np.zeros_like(input_edge) - 1 | |
| sub_mesh = mesh.subgraph(list(cc)).copy() | |
| ends_1, ends_2 = end_pts[0], end_pts[1] | |
| edge_id = global_mesh.nodes[ends_1]['edge_id'] | |
| npath = [*netx.shortest_path(sub_mesh, (ends_1[0], ends_1[1]), (ends_2[0], ends_2[1]), weight='length')] | |
| for np_node in npath: | |
| my_npath_map[np_node[0], np_node[1]] = edge_id | |
| fpath = [] | |
| if global_mesh.nodes[ends_1].get('far') is None: | |
| print("None far") | |
| else: | |
| fnodes = global_mesh.nodes[ends_1].get('far') | |
| dmask = mask + 0 | |
| while True: | |
| dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1) | |
| ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\ | |
| global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)] | |
| if len(ffnode) > 0: | |
| fnode = ffnode[0] | |
| break | |
| e_fnodes = global_mesh.nodes[ends_2].get('far') | |
| dmask = mask + 0 | |
| while True: | |
| dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1) | |
| e_ffnode = [e_fnode for e_fnode in e_fnodes if (dmask[e_fnode[0], e_fnode[1]] > 0 and mask[e_fnode[0], e_fnode[1]] == 0 and\ | |
| global_mesh.nodes[e_fnode].get('inpaint_id') != inpaint_id + 1)] | |
| if len(e_ffnode) > 0: | |
| e_fnode = e_ffnode[0] | |
| break | |
| fpath.append((fnode[0], fnode[1])) | |
| if len(e_ffnode) == 0 or len(ffnode) == 0: | |
| return my_npath_map, my_fpath_map, [], [] | |
| barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]]) | |
| n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1])) | |
| while True: | |
| if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]: | |
| n2f_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| for step in range(0, len(npath)): | |
| if step == 0: | |
| continue | |
| elif step == 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]: | |
| next_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| elif step > 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]: | |
| next_barrel = barrel_pair.copy() | |
| break | |
| barrel_pair = np.roll(barrel_pair, 1, axis=1) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| new_locs = [] | |
| if abs(n2f_dir[0]) == 1: | |
| new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1])) | |
| if abs(n2f_dir[1]) == 1: | |
| new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1])) | |
| if len(new_locs) > 1: | |
| new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1]))) | |
| break_flag = False | |
| for new_loc in new_locs: | |
| new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]), | |
| (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\ | |
| if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]] | |
| if fpath_map is not None and np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != 0: | |
| break_flag = True | |
| break | |
| if my_npath_map[new_loc[0], new_loc[1]] != -1: | |
| continue | |
| if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id: | |
| break_flag = True | |
| break | |
| fpath.append(new_loc) | |
| if break_flag is True: | |
| break | |
| if (e_fnode[0], e_fnode[1]) not in fpath: | |
| fpath.append((e_fnode[0], e_fnode[1])) | |
| if step != len(npath) - 1: | |
| for xx in npath[step:]: | |
| if my_npath_map[xx[0], xx[1]] == edge_id: | |
| my_npath_map[xx[0], xx[1]] = -1 | |
| npath = npath[:step] | |
| if len(fpath) > 0: | |
| for fp_node in fpath: | |
| my_fpath_map[fp_node[0], fp_node[1]] = edge_id | |
| return my_fpath_map, my_npath_map, npath, fpath | |
| def plan_path(mesh, info_on_pix, cc, end_pt, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=None): | |
| my_npath_map = np.zeros_like(input_edge) - 1 | |
| my_fpath_map = np.zeros_like(input_edge) - 1 | |
| sub_mesh = mesh.subgraph(list(cc)).copy() | |
| pnodes = netx.periphery(sub_mesh) | |
| ends = [*end_pt] | |
| edge_id = global_mesh.nodes[ends[0]]['edge_id'] | |
| pnodes = sorted(pnodes, | |
| key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])), | |
| reverse=True)[0] | |
| if npath is None: | |
| npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')] | |
| else: | |
| if (ends[0][0], ends[0][1]) == npath[0]: | |
| npath = npath | |
| elif (ends[0][0], ends[0][1]) == npath[-1]: | |
| npath = npath[::-1] | |
| else: | |
| import pdb; pdb.set_trace() | |
| for np_node in npath: | |
| my_npath_map[np_node[0], np_node[1]] = edge_id | |
| fpath = [] | |
| if global_mesh.nodes[ends[0]].get('far') is None: | |
| print("None far") | |
| else: | |
| fnodes = global_mesh.nodes[ends[0]].get('far') | |
| dmask = mask + 0 | |
| did = 0 | |
| while True: | |
| did += 1 | |
| if did > 3: | |
| return my_fpath_map, my_npath_map, -1 | |
| dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1) | |
| ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\ | |
| global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)] | |
| if len(ffnode) > 0: | |
| fnode = ffnode[0] | |
| break | |
| fpath.append((fnode[0], fnode[1])) | |
| disp_diff = 0. | |
| for n_loc in npath: | |
| if mask[n_loc[0], n_loc[1]] != 0: | |
| disp_diff = abs(abs(1. / info_on_pix[(n_loc[0], n_loc[1])][0]['depth']) - abs(1. / ends[0][2])) | |
| break | |
| barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]]) | |
| n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1])) | |
| while True: | |
| if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]: | |
| n2f_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| for step in range(0, len(npath)): | |
| if step == 0: | |
| continue | |
| elif step == 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]: | |
| next_barrel = barrel_dir.copy() | |
| break | |
| barrel_dir = np.roll(barrel_dir, 1, axis=0) | |
| barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| elif step > 1: | |
| next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1]) | |
| while True: | |
| if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]: | |
| next_barrel = barrel_pair.copy() | |
| break | |
| barrel_pair = np.roll(barrel_pair, 1, axis=1) | |
| n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1]) | |
| new_locs = [] | |
| if abs(n2f_dir[0]) == 1: | |
| new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1])) | |
| if abs(n2f_dir[1]) == 1: | |
| new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1])) | |
| if len(new_locs) > 1: | |
| new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1]))) | |
| break_flag = False | |
| for new_loc in new_locs: | |
| new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]), | |
| (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\ | |
| if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]] | |
| if fpath_map is not None and np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True: | |
| break_flag = True | |
| break | |
| if np.all([(my_fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True: | |
| break_flag = True | |
| break | |
| if my_npath_map[new_loc[0], new_loc[1]] != -1: | |
| continue | |
| if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id: | |
| break_flag = True | |
| break | |
| if valid_map[new_loc[0], new_loc[1]] == 0: | |
| break_flag = True | |
| break | |
| fpath.append(new_loc) | |
| if break_flag is True: | |
| break | |
| if step != len(npath) - 1: | |
| for xx in npath[step:]: | |
| if my_npath_map[xx[0], xx[1]] == edge_id: | |
| my_npath_map[xx[0], xx[1]] = -1 | |
| npath = npath[:step] | |
| if len(fpath) > 0: | |
| for fp_node in fpath: | |
| my_fpath_map[fp_node[0], fp_node[1]] = edge_id | |
| return my_fpath_map, my_npath_map, disp_diff | |
| def refresh_node(old_node, old_feat, new_node, new_feat, mesh, stime=False): | |
| mesh.add_node(new_node) | |
| mesh.nodes[new_node].update(new_feat) | |
| mesh.nodes[new_node].update(old_feat) | |
| for ne in mesh.neighbors(old_node): | |
| mesh.add_edge(new_node, ne) | |
| if mesh.nodes[new_node].get('far') is not None: | |
| tmp_far_nodes = mesh.nodes[new_node]['far'] | |
| for far_node in tmp_far_nodes: | |
| if mesh.has_node(far_node) is False: | |
| mesh.nodes[new_node]['far'].remove(far_node) | |
| continue | |
| if mesh.nodes[far_node].get('near') is not None: | |
| for idx in range(len(mesh.nodes[far_node].get('near'))): | |
| if mesh.nodes[far_node]['near'][idx][0] == new_node[0] and mesh.nodes[far_node]['near'][idx][1] == new_node[1]: | |
| if len(mesh.nodes[far_node]['near'][idx]) == len(old_node): | |
| mesh.nodes[far_node]['near'][idx] = new_node | |
| if mesh.nodes[new_node].get('near') is not None: | |
| tmp_near_nodes = mesh.nodes[new_node]['near'] | |
| for near_node in tmp_near_nodes: | |
| if mesh.has_node(near_node) is False: | |
| mesh.nodes[new_node]['near'].remove(near_node) | |
| continue | |
| if mesh.nodes[near_node].get('far') is not None: | |
| for idx in range(len(mesh.nodes[near_node].get('far'))): | |
| if mesh.nodes[near_node]['far'][idx][0] == new_node[0] and mesh.nodes[near_node]['far'][idx][1] == new_node[1]: | |
| if len(mesh.nodes[near_node]['far'][idx]) == len(old_node): | |
| mesh.nodes[near_node]['far'][idx] = new_node | |
| if new_node != old_node: | |
| mesh.remove_node(old_node) | |
| if stime is False: | |
| return mesh | |
| else: | |
| return mesh, None, None | |
| def create_placeholder(context, mask, depth, fpath_map, npath_map, mesh, inpaint_id, edge_ccs, extend_edge_cc, all_edge_maps, self_edge_id): | |
| add_node_time = 0 | |
| add_edge_time = 0 | |
| add_far_near_time = 0 | |
| valid_area = context + mask | |
| H, W = mesh.graph['H'], mesh.graph['W'] | |
| edge_cc = edge_ccs[self_edge_id] | |
| num_com = len(edge_cc) + len(extend_edge_cc) | |
| hxs, hys = np.where(mask > 0) | |
| for hx, hy in zip(hxs, hys): | |
| mesh.add_node((hx, hy), inpaint_id=inpaint_id + 1, num_context=num_com) | |
| for hx, hy in zip(hxs, hys): | |
| four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\ | |
| 0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and valid_area[x, y] != 0] | |
| for ne in four_nes: | |
| if mask[ne[0], ne[1]] != 0: | |
| if not mesh.has_edge((hx, hy), ne): | |
| mesh.add_edge((hx, hy), ne) | |
| elif depth[ne[0], ne[1]] != 0: | |
| if mesh.has_node((ne[0], ne[1], depth[ne[0], ne[1]])) and\ | |
| not mesh.has_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]])): | |
| mesh.add_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]])) | |
| else: | |
| print("Undefined context node.") | |
| import pdb; pdb.set_trace() | |
| near_ids = np.unique(npath_map) | |
| if near_ids[0] == -1: near_ids = near_ids[1:] | |
| for near_id in near_ids: | |
| hxs, hys = np.where((fpath_map == near_id) & (mask > 0)) | |
| if hxs.shape[0] > 0: | |
| mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1 | |
| else: | |
| break | |
| for hx, hy in zip(hxs, hys): | |
| mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id'])) | |
| four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\ | |
| x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and npath_map[x, y] == near_id] | |
| for xx in four_nes: | |
| xx_n = copy.deepcopy(xx) | |
| if not mesh.has_node(xx_n): | |
| if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])): | |
| xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]]) | |
| if mesh.has_edge((hx, hy), xx_n): | |
| # pass | |
| mesh.remove_edge((hx, hy), xx_n) | |
| if mesh.nodes[(hx, hy)].get('near') is None: | |
| mesh.nodes[(hx, hy)]['near'] = [] | |
| mesh.nodes[(hx, hy)]['near'].append(xx_n) | |
| connect_point_exception = set() | |
| hxs, hys = np.where((npath_map == near_id) & (all_edge_maps > -1)) | |
| for hx, hy in zip(hxs, hys): | |
| unknown_id = int(round(all_edge_maps[hx, hy])) | |
| if unknown_id != near_id and unknown_id != self_edge_id: | |
| unknown_node = set([xx for xx in edge_ccs[unknown_id] if xx[0] == hx and xx[1] == hy]) | |
| connect_point_exception |= unknown_node | |
| hxs, hys = np.where((npath_map == near_id) & (mask > 0)) | |
| if hxs.shape[0] > 0: | |
| mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1 | |
| else: | |
| break | |
| for hx, hy in zip(hxs, hys): | |
| mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id'])) | |
| mesh.nodes[(hx, hy)]['connect_point_id'] = int(round(near_id)) | |
| mesh.nodes[(hx, hy)]['connect_point_exception'] = connect_point_exception | |
| four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\ | |
| x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and fpath_map[x, y] == near_id] | |
| for xx in four_nes: | |
| xx_n = copy.deepcopy(xx) | |
| if not mesh.has_node(xx_n): | |
| if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])): | |
| xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]]) | |
| if mesh.has_edge((hx, hy), xx_n): | |
| mesh.remove_edge((hx, hy), xx_n) | |
| if mesh.nodes[(hx, hy)].get('far') is None: | |
| mesh.nodes[(hx, hy)]['far'] = [] | |
| mesh.nodes[(hx, hy)]['far'].append(xx_n) | |
| return mesh, add_node_time, add_edge_time, add_far_near_time | |
| def clean_far_edge(mask_edge, mask_edge_with_id, context_edge, mask, info_on_pix, global_mesh, anchor): | |
| if isinstance(mask_edge, torch.Tensor): | |
| if mask_edge.is_cuda: | |
| mask_edge = mask_edge.cpu() | |
| mask_edge = mask_edge.data | |
| mask_edge = mask_edge.numpy() | |
| if isinstance(context_edge, torch.Tensor): | |
| if context_edge.is_cuda: | |
| context_edge = context_edge.cpu() | |
| context_edge = context_edge.data | |
| context_edge = context_edge.numpy() | |
| if isinstance(mask, torch.Tensor): | |
| if mask.is_cuda: | |
| mask = mask.cpu() | |
| mask = mask.data | |
| mask = mask.numpy() | |
| mask = mask.squeeze() | |
| mask_edge = mask_edge.squeeze() | |
| context_edge = context_edge.squeeze() | |
| valid_near_edge = np.zeros_like(mask_edge) | |
| far_edge = np.zeros_like(mask_edge) | |
| far_edge_with_id = np.ones_like(mask_edge) * -1 | |
| near_edge_with_id = np.ones_like(mask_edge) * -1 | |
| uncleaned_far_edge = np.zeros_like(mask_edge) | |
| # Detect if there is any valid pixel mask_edge, if not ==> return default value | |
| if mask_edge.sum() == 0: | |
| return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id | |
| mask_edge_ids = dict(collections.Counter(mask_edge_with_id.flatten())).keys() | |
| for edge_id in mask_edge_ids: | |
| if edge_id < 0: | |
| continue | |
| specific_edge_map = (mask_edge_with_id == edge_id).astype(np.uint8) | |
| _, sub_specific_edge_maps = cv2.connectedComponents(specific_edge_map.astype(np.uint8), connectivity=8) | |
| for sub_edge_id in range(1, sub_specific_edge_maps.max() + 1): | |
| specific_edge_map = (sub_specific_edge_maps == sub_edge_id).astype(np.uint8) | |
| edge_pxs, edge_pys = np.where(specific_edge_map > 0) | |
| edge_mesh = netx.Graph() | |
| for edge_px, edge_py in zip(edge_pxs, edge_pys): | |
| edge_mesh.add_node((edge_px, edge_py)) | |
| for ex in [edge_px-1, edge_px, edge_px+1]: | |
| for ey in [edge_py-1, edge_py, edge_py+1]: | |
| if edge_px == ex and edge_py == ey: | |
| continue | |
| if ex < 0 or ex >= specific_edge_map.shape[0] or ey < 0 or ey >= specific_edge_map.shape[1]: | |
| continue | |
| if specific_edge_map[ex, ey] == 1: | |
| if edge_mesh.has_node((ex, ey)): | |
| edge_mesh.add_edge((ex, ey), (edge_px, edge_py)) | |
| periphery_nodes = netx.periphery(edge_mesh) | |
| path_diameter = netx.diameter(edge_mesh) | |
| start_near_node = None | |
| for node_s in periphery_nodes: | |
| for node_e in periphery_nodes: | |
| if node_s != node_e: | |
| if netx.shortest_path_length(edge_mesh, node_s, node_e) == path_diameter: | |
| if np.any(context_edge[node_s[0]-1:node_s[0]+2, node_s[1]-1:node_s[1]+2].flatten()): | |
| start_near_node = (node_s[0], node_s[1]) | |
| end_near_node = (node_e[0], node_e[1]) | |
| break | |
| if np.any(context_edge[node_e[0]-1:node_e[0]+2, node_e[1]-1:node_e[1]+2].flatten()): | |
| start_near_node = (node_e[0], node_e[1]) | |
| end_near_node = (node_s[0], node_s[1]) | |
| break | |
| if start_near_node is not None: | |
| break | |
| if start_near_node is None: | |
| continue | |
| new_specific_edge_map = np.zeros_like(mask) | |
| for path_node in netx.shortest_path(edge_mesh, start_near_node, end_near_node): | |
| new_specific_edge_map[path_node[0], path_node[1]] = 1 | |
| context_near_pxs, context_near_pys = np.where(context_edge[start_near_node[0]-1:start_near_node[0]+2, start_near_node[1]-1:start_near_node[1]+2] > 0) | |
| distance = np.abs((context_near_pxs - 1)) + np.abs((context_near_pys - 1)) | |
| if (np.where(distance == distance.min())[0].shape[0]) > 1: | |
| closest_pxs = context_near_pxs[np.where(distance == distance.min())[0]] | |
| closest_pys = context_near_pys[np.where(distance == distance.min())[0]] | |
| closest_depths = [] | |
| for closest_px, closest_py in zip(closest_pxs, closest_pys): | |
| if info_on_pix.get((closest_px + start_near_node[0] - 1 + anchor[0], closest_py + start_near_node[1] - 1 + anchor[2])) is not None: | |
| for info in info_on_pix.get((closest_px + start_near_node[0] - 1 + anchor[0], closest_py + start_near_node[1] - 1 + anchor[2])): | |
| if info['synthesis'] is False: | |
| closest_depths.append(abs(info['depth'])) | |
| context_near_px, context_near_py = closest_pxs[np.array(closest_depths).argmax()], closest_pys[np.array(closest_depths).argmax()] | |
| else: | |
| context_near_px, context_near_py = context_near_pxs[distance.argmin()], context_near_pys[distance.argmin()] | |
| context_near_node = (start_near_node[0]-1 + context_near_px, start_near_node[1]-1 + context_near_py) | |
| far_node_list = [] | |
| global_context_near_node = (context_near_node[0] + anchor[0], context_near_node[1] + anchor[2]) | |
| if info_on_pix.get(global_context_near_node) is not None: | |
| for info in info_on_pix[global_context_near_node]: | |
| if info['synthesis'] is False: | |
| context_near_node_3d = (global_context_near_node[0], global_context_near_node[1], info['depth']) | |
| if global_mesh.nodes[context_near_node_3d].get('far') is not None: | |
| for far_node in global_mesh.nodes[context_near_node_3d].get('far'): | |
| far_node = (far_node[0] - anchor[0], far_node[1] - anchor[2], far_node[2]) | |
| if mask[far_node[0], far_node[1]] == 0: | |
| far_node_list.append([far_node[0], far_node[1]]) | |
| if len(far_node_list) > 0: | |
| far_nodes_dist = np.sum(np.abs(np.array(far_node_list) - np.array([[edge_px, edge_py]])), axis=1) | |
| context_far_node = tuple(far_node_list[far_nodes_dist.argmin()]) | |
| corresponding_far_edge = np.zeros_like(mask_edge) | |
| corresponding_far_edge[context_far_node[0], context_far_node[1]] = 1 | |
| surround_map = cv2.dilate(new_specific_edge_map.astype(np.uint8), | |
| np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), | |
| iterations=1) | |
| specific_edge_map_wo_end_pt = new_specific_edge_map.copy() | |
| specific_edge_map_wo_end_pt[end_near_node[0], end_near_node[1]] = 0 | |
| surround_map_wo_end_pt = cv2.dilate(specific_edge_map_wo_end_pt.astype(np.uint8), | |
| np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), | |
| iterations=1) | |
| surround_map_wo_end_pt[new_specific_edge_map > 0] = 0 | |
| surround_map_wo_end_pt[context_near_node[0], context_near_node[1]] = 0 | |
| surround_map = surround_map_wo_end_pt.copy() | |
| _, far_edge_cc = cv2.connectedComponents(surround_map.astype(np.uint8), connectivity=4) | |
| start_far_node = None | |
| accompany_far_node = None | |
| if surround_map[context_far_node[0], context_far_node[1]] == 1: | |
| start_far_node = context_far_node | |
| else: | |
| four_nes = [(context_far_node[0] - 1, context_far_node[1]), | |
| (context_far_node[0] + 1, context_far_node[1]), | |
| (context_far_node[0], context_far_node[1] - 1), | |
| (context_far_node[0], context_far_node[1] + 1)] | |
| candidate_bevel = [] | |
| for ne in four_nes: | |
| if surround_map[ne[0], ne[1]] == 1: | |
| start_far_node = (ne[0], ne[1]) | |
| break | |
| elif (ne[0] != context_near_node[0] or ne[1] != context_near_node[1]) and \ | |
| (ne[0] != start_near_node[0] or ne[1] != start_near_node[1]): | |
| candidate_bevel.append((ne[0], ne[1])) | |
| if start_far_node is None: | |
| for ne in candidate_bevel: | |
| if ne[0] == context_far_node[0]: | |
| bevel_xys = [[ne[0] + 1, ne[1]], [ne[0] - 1, ne[1]]] | |
| if ne[1] == context_far_node[1]: | |
| bevel_xys = [[ne[0], ne[1] + 1], [ne[0], ne[1] - 1]] | |
| for bevel_x, bevel_y in bevel_xys: | |
| if surround_map[bevel_x, bevel_y] == 1: | |
| start_far_node = (bevel_x, bevel_y) | |
| accompany_far_node = (ne[0], ne[1]) | |
| break | |
| if start_far_node is not None: | |
| break | |
| if start_far_node is not None: | |
| for far_edge_id in range(1, far_edge_cc.max() + 1): | |
| specific_far_edge = (far_edge_cc == far_edge_id).astype(np.uint8) | |
| if specific_far_edge[start_far_node[0], start_far_node[1]] == 1: | |
| if accompany_far_node is not None: | |
| specific_far_edge[accompany_far_node] = 1 | |
| far_edge[specific_far_edge > 0] = 1 | |
| far_edge_with_id[specific_far_edge > 0] = edge_id | |
| end_far_candidates = np.zeros_like(far_edge) | |
| end_far_candidates[end_near_node[0], end_near_node[1]] = 1 | |
| end_far_candidates = cv2.dilate(end_far_candidates.astype(np.uint8), | |
| np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), | |
| iterations=1) | |
| end_far_candidates[end_near_node[0], end_near_node[1]] = 0 | |
| invalid_nodes = (((far_edge_cc != far_edge_id).astype(np.uint8) * \ | |
| (far_edge_cc != 0).astype(np.uint8)).astype(np.uint8) + \ | |
| (new_specific_edge_map).astype(np.uint8) + \ | |
| (mask == 0).astype(np.uint8)).clip(0, 1) | |
| end_far_candidates[invalid_nodes > 0] = 0 | |
| far_edge[end_far_candidates > 0] = 1 | |
| far_edge_with_id[end_far_candidates > 0] = edge_id | |
| far_edge[context_far_node[0], context_far_node[1]] = 1 | |
| far_edge_with_id[context_far_node[0], context_far_node[1]] = edge_id | |
| near_edge_with_id[(mask_edge_with_id == edge_id) > 0] = edge_id | |
| uncleaned_far_edge = far_edge.copy() | |
| far_edge[mask == 0] = 0 | |
| return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id | |
| def get_MiDaS_samples(image_folder, depth_folder, config, specific=None, aft_certain=None): | |
| lines = [os.path.splitext(os.path.basename(xx))[0] for xx in glob.glob(os.path.join(image_folder, '*' + config['img_format']))] | |
| samples = [] | |
| generic_pose = np.eye(4) | |
| assert len(config['traj_types']) == len(config['x_shift_range']) ==\ | |
| len(config['y_shift_range']) == len(config['z_shift_range']) == len(config['video_postfix']), \ | |
| "The number of elements in 'traj_types', 'x_shift_range', 'y_shift_range', 'z_shift_range' and \ | |
| 'video_postfix' should be equal." | |
| tgt_pose = [[generic_pose * 1]] | |
| tgts_poses = [] | |
| for traj_idx in range(len(config['traj_types'])): | |
| tgt_poses = [] | |
| sx, sy, sz = path_planning(config['num_frames'], config['x_shift_range'][traj_idx], config['y_shift_range'][traj_idx], | |
| config['z_shift_range'][traj_idx], path_type=config['traj_types'][traj_idx]) | |
| for xx, yy, zz in zip(sx, sy, sz): | |
| tgt_poses.append(generic_pose * 1.) | |
| tgt_poses[-1][:3, -1] = np.array([xx, yy, zz]) | |
| tgts_poses += [tgt_poses] | |
| tgt_pose = generic_pose * 1 | |
| aft_flag = True | |
| if aft_certain is not None and len(aft_certain) > 0: | |
| aft_flag = False | |
| for seq_dir in lines: | |
| if specific is not None and len(specific) > 0: | |
| if specific != seq_dir: | |
| continue | |
| if aft_certain is not None and len(aft_certain) > 0: | |
| if aft_certain == seq_dir: | |
| aft_flag = True | |
| if aft_flag is False: | |
| continue | |
| samples.append({}) | |
| sdict = samples[-1] | |
| sdict['depth_fi'] = os.path.join(depth_folder, seq_dir + config['depth_format']) | |
| sdict['ref_img_fi'] = os.path.join(image_folder, seq_dir + config['img_format']) | |
| H, W = imageio.imread(sdict['ref_img_fi']).shape[:2] | |
| sdict['int_mtx'] = np.array([[max(H, W), 0, W//2], [0, max(H, W), H//2], [0, 0, 1]]).astype(np.float32) | |
| if sdict['int_mtx'].max() > 1: | |
| sdict['int_mtx'][0, :] = sdict['int_mtx'][0, :] / float(W) | |
| sdict['int_mtx'][1, :] = sdict['int_mtx'][1, :] / float(H) | |
| sdict['ref_pose'] = np.eye(4) | |
| sdict['tgt_pose'] = tgt_pose | |
| sdict['tgts_poses'] = tgts_poses | |
| sdict['video_postfix'] = config['video_postfix'] | |
| sdict['tgt_name'] = [os.path.splitext(os.path.basename(sdict['depth_fi']))[0]] | |
| sdict['src_pair_name'] = sdict['tgt_name'][0] | |
| return samples | |
| def get_valid_size(imap): | |
| x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1 | |
| x_min = np.where(imap.sum(1).squeeze() > 0)[0].min() | |
| y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1 | |
| y_min = np.where(imap.sum(0).squeeze() > 0)[0].min() | |
| size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min} | |
| return size_dict | |
| def dilate_valid_size(isize_dict, imap, dilate=[0, 0]): | |
| osize_dict = copy.deepcopy(isize_dict) | |
| osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0]) | |
| osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0]) | |
| osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0]) | |
| osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1]) | |
| return osize_dict | |
| def crop_maps_by_size(size, *imaps): | |
| omaps = [] | |
| for imap in imaps: | |
| omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy()) | |
| return omaps | |
| def smooth_cntsyn_gap(init_depth_map, mask_region, context_region, init_mask_region=None): | |
| if init_mask_region is not None: | |
| curr_mask_region = init_mask_region * 1 | |
| else: | |
| curr_mask_region = mask_region * 0 | |
| depth_map = init_depth_map.copy() | |
| for _ in range(2): | |
| cm_mask = context_region + curr_mask_region | |
| depth_s1 = np.roll(depth_map, 1, 0) | |
| depth_s2 = np.roll(depth_map, -1, 0) | |
| depth_s3 = np.roll(depth_map, 1, 1) | |
| depth_s4 = np.roll(depth_map, -1, 1) | |
| mask_s1 = np.roll(cm_mask, 1, 0) | |
| mask_s2 = np.roll(cm_mask, -1, 0) | |
| mask_s3 = np.roll(cm_mask, 1, 1) | |
| mask_s4 = np.roll(cm_mask, -1, 1) | |
| fluxin_depths = (depth_s1 * mask_s1 + depth_s2 * mask_s2 + depth_s3 * mask_s3 + depth_s4 * mask_s4) / \ | |
| ((mask_s1 + mask_s2 + mask_s3 + mask_s4) + 1e-6) | |
| fluxin_mask = (fluxin_depths != 0) * mask_region | |
| init_mask = (fluxin_mask * (curr_mask_region >= 0).astype(np.float32) > 0).astype(np.uint8) | |
| depth_map[init_mask > 0] = fluxin_depths[init_mask > 0] | |
| if init_mask.shape[-1] > curr_mask_region.shape[-1]: | |
| curr_mask_region[init_mask.sum(-1, keepdims=True) > 0] = 1 | |
| else: | |
| curr_mask_region[init_mask > 0] = 1 | |
| depth_map[fluxin_mask > 0] = fluxin_depths[fluxin_mask > 0] | |
| return depth_map | |
| def read_MiDaS_depth(disp_fi, disp_rescale=10., h=None, w=None): | |
| if 'npy' in os.path.splitext(disp_fi)[-1]: | |
| disp = np.load(disp_fi) | |
| else: | |
| disp = imageio.imread(disp_fi).astype(np.float32) | |
| disp = disp - disp.min() | |
| disp = cv2.blur(disp / disp.max(), ksize=(3, 3)) * disp.max() | |
| disp = (disp / disp.max()) * disp_rescale | |
| if h is not None and w is not None: | |
| disp = resize(disp / disp.max(), (h, w), order=1) * disp.max() | |
| depth = 1. / np.maximum(disp, 0.05) | |
| return depth | |
| def follow_image_aspect_ratio(depth, image): | |
| H, W = image.shape[:2] | |
| image_aspect_ratio = H / W | |
| dH, dW = depth.shape[:2] | |
| depth_aspect_ratio = dH / dW | |
| if depth_aspect_ratio > image_aspect_ratio: | |
| resize_H = dH | |
| resize_W = dH / image_aspect_ratio | |
| else: | |
| resize_W = dW | |
| resize_H = dW * image_aspect_ratio | |
| depth = resize(depth / depth.max(), | |
| (int(resize_H), | |
| int(resize_W)), | |
| order=0) * depth.max() | |
| return depth | |
| def depth_resize(depth, origin_size, image_size): | |
| if origin_size[0] is not 0: | |
| max_depth = depth.max() | |
| depth = depth / max_depth | |
| depth = resize(depth, origin_size, order=1, mode='edge') | |
| depth = depth * max_depth | |
| else: | |
| max_depth = depth.max() | |
| depth = depth / max_depth | |
| depth = resize(depth, image_size, order=1, mode='edge') | |
| depth = depth * max_depth | |
| return depth | |
| def filter_irrelevant_edge(self_edge, other_edges, other_edges_with_id, current_edge_id, context, edge_ccs, mesh, anchor): | |
| other_edges = other_edges.squeeze() | |
| other_edges_with_id = other_edges_with_id.squeeze() | |
| self_edge = self_edge.squeeze() | |
| dilate_self_edge = cv2.dilate(self_edge.astype(np.uint8), np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), iterations=1) | |
| edge_ids = collections.Counter(other_edges_with_id.flatten()).keys() | |
| other_edges_info = [] | |
| # import ipdb | |
| # ipdb.set_trace() | |
| for edge_id in edge_ids: | |
| edge_id = int(edge_id) | |
| if edge_id >= 0: | |
| condition = ((other_edges_with_id == edge_id) * other_edges * context).astype(np.uint8) | |
| if dilate_self_edge[condition > 0].sum() == 0: | |
| other_edges[other_edges_with_id == edge_id] = 0 | |
| else: | |
| num_condition, condition_labels = cv2.connectedComponents(condition, connectivity=8) | |
| for condition_id in range(1, num_condition): | |
| isolate_condition = ((condition_labels == condition_id) > 0).astype(np.uint8) | |
| num_end_group, end_group = cv2.connectedComponents(((dilate_self_edge * isolate_condition) > 0).astype(np.uint8), connectivity=8) | |
| if num_end_group == 1: | |
| continue | |
| for end_id in range(1, num_end_group): | |
| end_pxs, end_pys = np.where((end_group == end_id)) | |
| end_px, end_py = end_pxs[0], end_pys[0] | |
| other_edges_info.append({}) | |
| other_edges_info[-1]['edge_id'] = edge_id | |
| # other_edges_info[-1]['near_depth'] = None | |
| other_edges_info[-1]['diff'] = None | |
| other_edges_info[-1]['edge_map'] = np.zeros_like(self_edge) | |
| other_edges_info[-1]['end_point_map'] = np.zeros_like(self_edge) | |
| other_edges_info[-1]['end_point_map'][(end_group == end_id)] = 1 | |
| other_edges_info[-1]['forbidden_point_map'] = np.zeros_like(self_edge) | |
| other_edges_info[-1]['forbidden_point_map'][(end_group != end_id) * (end_group != 0)] = 1 | |
| other_edges_info[-1]['forbidden_point_map'] = cv2.dilate(other_edges_info[-1]['forbidden_point_map'], kernel=np.array([[1,1,1],[1,1,1],[1,1,1]]), iterations=2) | |
| for x in edge_ccs[edge_id]: | |
| nx = x[0] - anchor[0] | |
| ny = x[1] - anchor[1] | |
| if nx == end_px and ny == end_py: | |
| # other_edges_info[-1]['near_depth'] = abs(nx) | |
| if mesh.nodes[x].get('far') is not None and len(mesh.nodes[x].get('far')) == 1: | |
| other_edges_info[-1]['diff'] = abs(1./abs([*mesh.nodes[x].get('far')][0][2]) - 1./abs(x[2])) | |
| else: | |
| other_edges_info[-1]['diff'] = 0 | |
| # if end_group[nx, ny] != end_id and end_group[nx, ny] > 0: | |
| # continue | |
| try: | |
| if isolate_condition[nx, ny] == 1: | |
| other_edges_info[-1]['edge_map'][nx, ny] = 1 | |
| except: | |
| pass | |
| try: | |
| other_edges_info = sorted(other_edges_info, key=lambda x : x['diff'], reverse=True) | |
| except: | |
| import pdb | |
| pdb.set_trace() | |
| # import pdb | |
| # pdb.set_trace() | |
| # other_edges = other_edges[..., None] | |
| for other_edge in other_edges_info: | |
| if other_edge['end_point_map'] is None: | |
| import pdb | |
| pdb.set_trace() | |
| other_edges = other_edges * context | |
| return other_edges, other_edges_info | |
| def require_depth_edge(context_edge, mask): | |
| dilate_mask = cv2.dilate(mask, np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), iterations=1) | |
| if (dilate_mask * context_edge).max() == 0: | |
| return False | |
| else: | |
| return True | |
| def refine_color_around_edge(mesh, info_on_pix, edge_ccs, config, spdb=False): | |
| H, W = mesh.graph['H'], mesh.graph['W'] | |
| tmp_edge_ccs = copy.deepcopy(edge_ccs) | |
| for edge_id, edge_cc in enumerate(edge_ccs): | |
| if len(edge_cc) == 0: | |
| continue | |
| near_maps = np.zeros((H, W)).astype(np.bool) | |
| far_maps = np.zeros((H, W)).astype(np.bool) | |
| tmp_far_nodes = set() | |
| far_nodes = set() | |
| near_nodes = set() | |
| end_nodes = set() | |
| for i in range(5): | |
| if i == 0: | |
| for edge_node in edge_cc: | |
| if mesh.nodes[edge_node].get('depth_edge_dilate_2_color_flag') is not True: | |
| break | |
| if mesh.nodes[edge_node].get('inpaint_id') == 1: | |
| near_nodes.add(edge_node) | |
| tmp_node = mesh.nodes[edge_node].get('far') | |
| tmp_node = set(tmp_node) if tmp_node is not None else set() | |
| tmp_far_nodes |= tmp_node | |
| rmv_tmp_far_nodes = set() | |
| for far_node in tmp_far_nodes: | |
| if not(mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1): | |
| rmv_tmp_far_nodes.add(far_node) | |
| if len(tmp_far_nodes - rmv_tmp_far_nodes) == 0: | |
| break | |
| else: | |
| for near_node in near_nodes: | |
| near_maps[near_node[0], near_node[1]] = True | |
| mesh.nodes[near_node]['refine_rgbd'] = True | |
| mesh.nodes[near_node]['backup_depth'] = near_node[2] \ | |
| if mesh.nodes[near_node].get('real_depth') is None else mesh.nodes[near_node]['real_depth'] | |
| mesh.nodes[near_node]['backup_color'] = mesh.nodes[near_node]['color'] | |
| for far_node in tmp_far_nodes: | |
| if mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1: | |
| far_nodes.add(far_node) | |
| far_maps[far_node[0], far_node[1]] = True | |
| mesh.nodes[far_node]['refine_rgbd'] = True | |
| mesh.nodes[far_node]['backup_depth'] = far_node[2] \ | |
| if mesh.nodes[far_node].get('real_depth') is None else mesh.nodes[far_node]['real_depth'] | |
| mesh.nodes[far_node]['backup_color'] = mesh.nodes[far_node]['color'] | |
| tmp_far_nodes = far_nodes | |
| tmp_near_nodes = near_nodes | |
| else: | |
| tmp_far_nodes = new_tmp_far_nodes | |
| tmp_near_nodes = new_tmp_near_nodes | |
| new_tmp_far_nodes = None | |
| new_tmp_near_nodes = None | |
| new_tmp_far_nodes = set() | |
| new_tmp_near_nodes = set() | |
| for node in tmp_near_nodes: | |
| for ne_node in mesh.neighbors(node): | |
| if far_maps[ne_node[0], ne_node[1]] == False and \ | |
| near_maps[ne_node[0], ne_node[1]] == False: | |
| if mesh.nodes[ne_node].get('inpaint_id') == 1: | |
| new_tmp_near_nodes.add(ne_node) | |
| near_maps[ne_node[0], ne_node[1]] = True | |
| mesh.nodes[ne_node]['refine_rgbd'] = True | |
| mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \ | |
| if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth'] | |
| mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color'] | |
| else: | |
| mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \ | |
| if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth'] | |
| mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color'] | |
| end_nodes.add(node) | |
| near_nodes.update(new_tmp_near_nodes) | |
| for node in tmp_far_nodes: | |
| for ne_node in mesh.neighbors(node): | |
| if far_maps[ne_node[0], ne_node[1]] == False and \ | |
| near_maps[ne_node[0], ne_node[1]] == False: | |
| if mesh.nodes[ne_node].get('inpaint_id') == 1: | |
| new_tmp_far_nodes.add(ne_node) | |
| far_maps[ne_node[0], ne_node[1]] = True | |
| mesh.nodes[ne_node]['refine_rgbd'] = True | |
| mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \ | |
| if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth'] | |
| mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color'] | |
| else: | |
| mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \ | |
| if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth'] | |
| mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color'] | |
| end_nodes.add(node) | |
| far_nodes.update(new_tmp_far_nodes) | |
| if len(far_nodes) == 0: | |
| tmp_edge_ccs[edge_id] = set() | |
| continue | |
| for node in new_tmp_far_nodes | new_tmp_near_nodes: | |
| for ne_node in mesh.neighbors(node): | |
| if far_maps[ne_node[0], ne_node[1]] == False and near_maps[ne_node[0], ne_node[1]] == False: | |
| end_nodes.add(node) | |
| mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \ | |
| if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth'] | |
| mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color'] | |
| tmp_end_nodes = end_nodes | |
| refine_nodes = near_nodes | far_nodes | |
| remain_refine_nodes = copy.deepcopy(refine_nodes) | |
| accum_idx = 0 | |
| while len(remain_refine_nodes) > 0: | |
| accum_idx += 1 | |
| if accum_idx > 100: | |
| break | |
| new_tmp_end_nodes = None | |
| new_tmp_end_nodes = set() | |
| survive_tmp_end_nodes = set() | |
| for node in tmp_end_nodes: | |
| re_depth, re_color, re_count = 0, np.array([0., 0., 0.]), 0 | |
| for ne_node in mesh.neighbors(node): | |
| if mesh.nodes[ne_node].get('refine_rgbd') is True: | |
| if ne_node not in tmp_end_nodes: | |
| new_tmp_end_nodes.add(ne_node) | |
| else: | |
| try: | |
| re_depth += mesh.nodes[ne_node]['backup_depth'] | |
| re_color += mesh.nodes[ne_node]['backup_color'].astype(np.float32) | |
| re_count += 1. | |
| except: | |
| import pdb; pdb.set_trace() | |
| if re_count > 0: | |
| re_depth = re_depth / re_count | |
| re_color = re_color / re_count | |
| mesh.nodes[node]['backup_depth'] = re_depth | |
| mesh.nodes[node]['backup_color'] = re_color | |
| mesh.nodes[node]['refine_rgbd'] = False | |
| else: | |
| survive_tmp_end_nodes.add(node) | |
| for node in tmp_end_nodes - survive_tmp_end_nodes: | |
| if node in remain_refine_nodes: | |
| remain_refine_nodes.remove(node) | |
| tmp_end_nodes = new_tmp_end_nodes | |
| if spdb == True: | |
| bfrd_canvas = np.zeros((H, W)) | |
| bfrc_canvas = np.zeros((H, W, 3)).astype(np.uint8) | |
| aftd_canvas = np.zeros((H, W)) | |
| aftc_canvas = np.zeros((H, W, 3)).astype(np.uint8) | |
| for node in refine_nodes: | |
| bfrd_canvas[node[0], node[1]] = abs(node[2]) | |
| aftd_canvas[node[0], node[1]] = abs(mesh.nodes[node]['backup_depth']) | |
| bfrc_canvas[node[0], node[1]] = mesh.nodes[node]['color'].astype(np.uint8) | |
| aftc_canvas[node[0], node[1]] = mesh.nodes[node]['backup_color'].astype(np.uint8) | |
| f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharex=True, sharey=True); | |
| ax1.imshow(bfrd_canvas); | |
| ax2.imshow(aftd_canvas); | |
| ax3.imshow(bfrc_canvas); | |
| ax4.imshow(aftc_canvas); | |
| plt.show() | |
| import pdb; pdb.set_trace() | |
| for node in refine_nodes: | |
| if mesh.nodes[node].get('refine_rgbd') is not None: | |
| mesh.nodes[node].pop('refine_rgbd') | |
| mesh.nodes[node]['color'] = mesh.nodes[node]['backup_color'] | |
| for info in info_on_pix[(node[0], node[1])]: | |
| if info['depth'] == node[2]: | |
| info['color'] = mesh.nodes[node]['backup_color'] | |
| return mesh, info_on_pix | |
| def refine_depth_around_edge(mask_depth, far_edge, uncleaned_far_edge, near_edge, mask, all_depth, config): | |
| if isinstance(mask_depth, torch.Tensor): | |
| if mask_depth.is_cuda: | |
| mask_depth = mask_depth.cpu() | |
| mask_depth = mask_depth.data | |
| mask_depth = mask_depth.numpy() | |
| if isinstance(far_edge, torch.Tensor): | |
| if far_edge.is_cuda: | |
| far_edge = far_edge.cpu() | |
| far_edge = far_edge.data | |
| far_edge = far_edge.numpy() | |
| if isinstance(uncleaned_far_edge, torch.Tensor): | |
| if uncleaned_far_edge.is_cuda: | |
| uncleaned_far_edge = uncleaned_far_edge.cpu() | |
| uncleaned_far_edge = uncleaned_far_edge.data | |
| uncleaned_far_edge = uncleaned_far_edge.numpy() | |
| if isinstance(near_edge, torch.Tensor): | |
| if near_edge.is_cuda: | |
| near_edge = near_edge.cpu() | |
| near_edge = near_edge.data | |
| near_edge = near_edge.numpy() | |
| if isinstance(mask, torch.Tensor): | |
| if mask.is_cuda: | |
| mask = mask.cpu() | |
| mask = mask.data | |
| mask = mask.numpy() | |
| mask = mask.squeeze() | |
| uncleaned_far_edge = uncleaned_far_edge.squeeze() | |
| far_edge = far_edge.squeeze() | |
| near_edge = near_edge.squeeze() | |
| mask_depth = mask_depth.squeeze() | |
| dilate_far_edge = cv2.dilate(uncleaned_far_edge.astype(np.uint8), kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1) | |
| near_edge[dilate_far_edge == 0] = 0 | |
| dilate_near_edge = cv2.dilate(near_edge.astype(np.uint8), kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1) | |
| far_edge[dilate_near_edge == 0] = 0 | |
| init_far_edge = far_edge.copy() | |
| init_near_edge = near_edge.copy() | |
| for i in range(config['depth_edge_dilate_2']): | |
| init_far_edge = cv2.dilate(init_far_edge, kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1) | |
| init_far_edge[init_near_edge == 1] = 0 | |
| init_near_edge = cv2.dilate(init_near_edge, kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1) | |
| init_near_edge[init_far_edge == 1] = 0 | |
| init_far_edge[mask == 0] = 0 | |
| init_near_edge[mask == 0] = 0 | |
| hole_far_edge = 1 - init_far_edge | |
| hole_near_edge = 1 - init_near_edge | |
| change = None | |
| while True: | |
| change = False | |
| hole_far_edge[init_near_edge == 1] = 0 | |
| hole_near_edge[init_far_edge == 1] = 0 | |
| far_pxs, far_pys = np.where((hole_far_edge == 0) * (init_far_edge == 1) > 0) | |
| current_hole_far_edge = hole_far_edge.copy() | |
| for far_px, far_py in zip(far_pxs, far_pys): | |
| min_px = max(far_px - 1, 0) | |
| max_px = min(far_px + 2, mask.shape[0]-1) | |
| min_py = max(far_py - 1, 0) | |
| max_py = min(far_py + 2, mask.shape[1]-1) | |
| hole_far = current_hole_far_edge[min_px: max_px, min_py: max_py] | |
| tmp_mask = mask[min_px: max_px, min_py: max_py] | |
| all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0 | |
| all_depth_mask = (all_depth_patch != 0).astype(np.uint8) | |
| cross_element = np.array([[0,1,0],[1,1,1],[0,1,0]])[min_px - (far_px - 1): max_px - (far_px - 1), min_py - (far_py - 1): max_py - (far_py - 1)] | |
| combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_far * cross_element | |
| tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch) | |
| number = np.count_nonzero(tmp_patch) | |
| if number > 0: | |
| mask_depth[far_px, far_py] = np.sum(tmp_patch).astype(np.float32) / max(number, 1e-6) | |
| hole_far_edge[far_px, far_py] = 1 | |
| change = True | |
| near_pxs, near_pys = np.where((hole_near_edge == 0) * (init_near_edge == 1) > 0) | |
| current_hole_near_edge = hole_near_edge.copy() | |
| for near_px, near_py in zip(near_pxs, near_pys): | |
| min_px = max(near_px - 1, 0) | |
| max_px = min(near_px + 2, mask.shape[0]-1) | |
| min_py = max(near_py - 1, 0) | |
| max_py = min(near_py + 2, mask.shape[1]-1) | |
| hole_near = current_hole_near_edge[min_px: max_px, min_py: max_py] | |
| tmp_mask = mask[min_px: max_px, min_py: max_py] | |
| all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0 | |
| all_depth_mask = (all_depth_patch != 0).astype(np.uint8) | |
| cross_element = np.array([[0,1,0],[1,1,1],[0,1,0]])[min_px - near_px + 1:max_px - near_px + 1, min_py - near_py + 1:max_py - near_py + 1] | |
| combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_near * cross_element | |
| tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch) | |
| number = np.count_nonzero(tmp_patch) | |
| if number > 0: | |
| mask_depth[near_px, near_py] = np.sum(tmp_patch) / max(number, 1e-6) | |
| hole_near_edge[near_px, near_py] = 1 | |
| change = True | |
| if change is False: | |
| break | |
| return mask_depth | |
| def vis_depth_edge_connectivity(depth, config): | |
| disp = 1./depth | |
| u_diff = (disp[1:, :] - disp[:-1, :])[:-1, 1:-1] | |
| b_diff = (disp[:-1, :] - disp[1:, :])[1:, 1:-1] | |
| l_diff = (disp[:, 1:] - disp[:, :-1])[1:-1, :-1] | |
| r_diff = (disp[:, :-1] - disp[:, 1:])[1:-1, 1:] | |
| u_over = (np.abs(u_diff) > config['depth_threshold']).astype(np.float32) | |
| b_over = (np.abs(b_diff) > config['depth_threshold']).astype(np.float32) | |
| l_over = (np.abs(l_diff) > config['depth_threshold']).astype(np.float32) | |
| r_over = (np.abs(r_diff) > config['depth_threshold']).astype(np.float32) | |
| concat_diff = np.stack([u_diff, b_diff, r_diff, l_diff], axis=-1) | |
| concat_over = np.stack([u_over, b_over, r_over, l_over], axis=-1) | |
| over_diff = concat_diff * concat_over | |
| pos_over = (over_diff > 0).astype(np.float32).sum(-1).clip(0, 1) | |
| neg_over = (over_diff < 0).astype(np.float32).sum(-1).clip(0, 1) | |
| neg_over[(over_diff > 0).astype(np.float32).sum(-1) > 0] = 0 | |
| _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8) | |
| T_junction_maps = np.zeros_like(pos_over) | |
| for edge_id in range(1, edge_label.max() + 1): | |
| edge_map = (edge_label == edge_id).astype(np.uint8) | |
| edge_map = np.pad(edge_map, pad_width=((1,1),(1,1)), mode='constant') | |
| four_direc = np.roll(edge_map, 1, 1) + np.roll(edge_map, -1, 1) + np.roll(edge_map, 1, 0) + np.roll(edge_map, -1, 0) | |
| eight_direc = np.roll(np.roll(edge_map, 1, 1), 1, 0) + np.roll(np.roll(edge_map, 1, 1), -1, 0) + \ | |
| np.roll(np.roll(edge_map, -1, 1), 1, 0) + np.roll(np.roll(edge_map, -1, 1), -1, 0) | |
| eight_direc = (eight_direc + four_direc)[1:-1,1:-1] | |
| pos_over[eight_direc > 2] = 0 | |
| T_junction_maps[eight_direc > 2] = 1 | |
| _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8) | |
| edge_label = np.pad(edge_label, 1, mode='constant') | |
| return edge_label | |
| def max_size(mat, value=0): | |
| if not (mat and mat[0]): return (0, 0) | |
| it = iter(mat) | |
| prev = [(el==value) for el in next(it)] | |
| max_size = max_rectangle_size(prev) | |
| for row in it: | |
| hist = [(1+h) if el == value else 0 for h, el in zip(prev, row)] | |
| max_size = max(max_size, max_rectangle_size(hist), key=get_area) | |
| prev = hist | |
| return max_size | |
| def max_rectangle_size(histogram): | |
| Info = namedtuple('Info', 'start height') | |
| stack = [] | |
| top = lambda: stack[-1] | |
| max_size = (0, 0) # height, width of the largest rectangle | |
| pos = 0 # current position in the histogram | |
| for pos, height in enumerate(histogram): | |
| start = pos # position where rectangle starts | |
| while True: | |
| if not stack or height > top().height: | |
| stack.append(Info(start, height)) # push | |
| if stack and height < top().height: | |
| max_size = max(max_size, (top().height, (pos-top().start)), | |
| key=get_area) | |
| start, _ = stack.pop() | |
| continue | |
| break # height == top().height goes here | |
| pos += 1 | |
| for start, height in stack: | |
| max_size = max(max_size, (height, (pos-start)), | |
| key=get_area) | |
| return max_size | |
| def get_area(size): | |
| return reduce(mul, size) | |
| def find_anchors(matrix): | |
| matrix = [[*x] for x in matrix] | |
| mh, mw = max_size(matrix) | |
| matrix = np.array(matrix) | |
| # element = np.zeros((mh, mw)) | |
| for i in range(matrix.shape[0] + 1 - mh): | |
| for j in range(matrix.shape[1] + 1 - mw): | |
| if matrix[i:i + mh, j:j + mw].max() == 0: | |
| return i, i + mh, j, j + mw | |
| def find_largest_rect(dst_img, bg_color=(128, 128, 128)): | |
| valid = np.any(dst_img[..., :3] != bg_color, axis=-1) | |
| dst_h, dst_w = dst_img.shape[:2] | |
| ret, labels = cv2.connectedComponents(np.uint8(valid == False)) | |
| red_mat = np.zeros_like(labels) | |
| # denoise | |
| for i in range(1, np.max(labels)+1, 1): | |
| x, y, w, h = cv2.boundingRect(np.uint8(labels==i)) | |
| if x == 0 or (x+w) == dst_h or y == 0 or (y+h) == dst_w: | |
| red_mat[labels==i] = 1 | |
| # crop | |
| t, b, l, r = find_anchors(red_mat) | |
| return t, b, l, r | |