Spaces:
Runtime error
Runtime error
import os | |
import numpy as np | |
try: | |
import cynetworkx as netx | |
except ImportError: | |
import networkx as netx | |
import matplotlib.pyplot as plt | |
from functools import partial | |
from vispy import scene, io | |
from vispy.scene import visuals | |
from vispy.visuals.filters import Alpha | |
import cv2 | |
from moviepy.editor import ImageSequenceClip | |
from skimage.transform import resize | |
import time | |
import copy | |
import torch | |
import os | |
from utils import path_planning, open_small_mask, clean_far_edge, refine_depth_around_edge | |
from utils import refine_color_around_edge, filter_irrelevant_edge_new, require_depth_edge, clean_far_edge_new | |
from utils import create_placeholder, refresh_node, find_largest_rect | |
from mesh_tools import get_depth_from_maps, get_map_from_ccs, get_edge_from_nodes, get_depth_from_nodes, get_rgb_from_nodes, crop_maps_by_size, convert2tensor, recursive_add_edge, update_info, filter_edge, relabel_node, depth_inpainting | |
from mesh_tools import refresh_bord_depth, enlarge_border, fill_dummy_bord, extrapolate, fill_missing_node, incomplete_node, get_valid_size, dilate_valid_size, size_operation | |
import transforms3d | |
import random | |
from functools import reduce | |
def create_mesh(depth, image, int_mtx, config): | |
H, W, C = image.shape | |
ext_H, ext_W = H + 2 * config['extrapolation_thickness'], W + 2 * config['extrapolation_thickness'] | |
LDI = netx.Graph(H=ext_H, W=ext_W, noext_H=H, noext_W=W, cam_param=int_mtx) | |
xy2depth = {} | |
int_mtx_pix = int_mtx * np.array([[W], [H], [1.]]) | |
LDI.graph['cam_param_pix'], LDI.graph['cam_param_pix_inv'] = int_mtx_pix, np.linalg.inv(int_mtx_pix) | |
disp = 1. / (-depth) | |
LDI.graph['hoffset'], LDI.graph['woffset'] = config['extrapolation_thickness'], config['extrapolation_thickness'] | |
LDI.graph['bord_up'], LDI.graph['bord_down'] = LDI.graph['hoffset'] + 0, LDI.graph['hoffset'] + H | |
LDI.graph['bord_left'], LDI.graph['bord_right'] = LDI.graph['woffset'] + 0, LDI.graph['woffset'] + W | |
for idx in range(H): | |
for idy in range(W): | |
x, y = idx + LDI.graph['hoffset'], idy + LDI.graph['woffset'] | |
LDI.add_node((x, y, -depth[idx, idy]), | |
color=image[idx, idy], | |
disp=disp[idx, idy], | |
synthesis=False, | |
cc_id=set()) | |
xy2depth[(x, y)] = [-depth[idx, idy]] | |
for x, y, d in LDI.nodes: | |
two_nes = [ne for ne in [(x+1, y), (x, y+1)] if ne[0] < LDI.graph['bord_down'] and ne[1] < LDI.graph['bord_right']] | |
[LDI.add_edge((ne[0], ne[1], xy2depth[ne][0]), (x, y, d)) for ne in two_nes] | |
LDI = calculate_fov(LDI) | |
image = np.pad(image, | |
pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']), | |
(config['extrapolation_thickness'], config['extrapolation_thickness']), | |
(0, 0)), | |
mode='constant') | |
depth = np.pad(depth, | |
pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']), | |
(config['extrapolation_thickness'], config['extrapolation_thickness'])), | |
mode='constant') | |
return LDI, xy2depth, image, depth | |
def tear_edges(mesh, threshold = 0.00025, xy2depth=None): | |
remove_edge_list = [] | |
remove_horizon, remove_vertical = np.zeros((2, mesh.graph['H'], mesh.graph['W'])) | |
mesh_nodes = mesh.nodes | |
for edge in mesh.edges: | |
if abs(mesh_nodes[edge[0]]['disp'] - mesh_nodes[edge[1]]['disp']) > threshold: | |
remove_edge_list.append((edge[0], edge[1])) | |
near, far = edge if abs(edge[0][2]) < abs(edge[1][2]) else edge[::-1] | |
mesh_nodes[far]['near'] = [] if mesh_nodes[far].get('near') is None else mesh_nodes[far]['near'].append(near) | |
mesh_nodes[near]['far'] = [] if mesh_nodes[near].get('far') is None else mesh_nodes[near]['far'].append(far) | |
if near[0] == far[0]: | |
remove_horizon[near[0], np.minimum(near[1], far[1])] = 1 | |
elif near[1] == far[1]: | |
remove_vertical[np.minimum(near[0], far[0]), near[1]] = 1 | |
mesh.remove_edges_from(remove_edge_list) | |
remove_edge_list = [] | |
dang_horizon = np.where(np.roll(remove_horizon, 1, 0) + np.roll(remove_horizon, -1, 0) - remove_horizon == 2) | |
dang_vertical = np.where(np.roll(remove_vertical, 1, 1) + np.roll(remove_vertical, -1, 1) - remove_vertical == 2) | |
horizon_condition = lambda x, y: mesh.graph['bord_up'] + 1 <= x < mesh.graph['bord_down'] - 1 | |
vertical_condition = lambda x, y: mesh.graph['bord_left'] + 1 <= y < mesh.graph['bord_right'] - 1 | |
prjto3d = lambda x, y: (x, y, xy2depth[(x, y)][0]) | |
node_existence = lambda x, y: mesh.has_node(prjto3d(x, y)) | |
for x, y in zip(dang_horizon[0], dang_horizon[1]): | |
if horizon_condition(x, y) and node_existence(x, y) and node_existence(x, y+1): | |
remove_edge_list.append((prjto3d(x, y), prjto3d(x, y+1))) | |
for x, y in zip(dang_vertical[0], dang_vertical[1]): | |
if vertical_condition(x, y) and node_existence(x, y) and node_existence(x+1, y): | |
remove_edge_list.append((prjto3d(x, y), prjto3d(x+1, y))) | |
mesh.remove_edges_from(remove_edge_list) | |
return mesh | |
def calculate_fov(mesh): | |
k = mesh.graph['cam_param'] | |
mesh.graph['hFov'] = 2 * np.arctan(1. / (2*k[0, 0])) | |
mesh.graph['vFov'] = 2 * np.arctan(1. / (2*k[1, 1])) | |
mesh.graph['aspect'] = mesh.graph['noext_H'] / mesh.graph['noext_W'] | |
return mesh | |
def calculate_fov_FB(mesh): | |
mesh.graph['aspect'] = mesh.graph['H'] / mesh.graph['W'] | |
if mesh.graph['H'] > mesh.graph['W']: | |
mesh.graph['hFov'] = 0.508015513 | |
half_short = np.tan(mesh.graph['hFov']/2.0) | |
half_long = half_short * mesh.graph['aspect'] | |
mesh.graph['vFov'] = 2.0 * np.arctan(half_long) | |
else: | |
mesh.graph['vFov'] = 0.508015513 | |
half_short = np.tan(mesh.graph['vFov']/2.0) | |
half_long = half_short / mesh.graph['aspect'] | |
mesh.graph['hFov'] = 2.0 * np.arctan(half_long) | |
return mesh | |
def reproject_3d_int_detail(sx, sy, z, k_00, k_02, k_11, k_12, w_offset, h_offset): | |
abs_z = abs(z) | |
return [abs_z * ((sy+0.5-w_offset) * k_00 + k_02), abs_z * ((sx+0.5-h_offset) * k_11 + k_12), abs_z] | |
def reproject_3d_int_detail_FB(sx, sy, z, w_offset, h_offset, mesh): | |
if mesh.graph.get('tan_hFov') is None: | |
mesh.graph['tan_hFov'] = np.tan(mesh.graph['hFov'] / 2.) | |
if mesh.graph.get('tan_vFov') is None: | |
mesh.graph['tan_vFov'] = np.tan(mesh.graph['vFov'] / 2.) | |
ray = np.array([(-1. + 2. * ((sy+0.5-w_offset)/(mesh.graph['W'] - 1))) * mesh.graph['tan_hFov'], | |
(1. - 2. * (sx+0.5-h_offset)/(mesh.graph['H'] - 1)) * mesh.graph['tan_vFov'], | |
-1]) | |
point_3d = ray * np.abs(z) | |
return point_3d | |
def reproject_3d_int(sx, sy, z, mesh): | |
k = mesh.graph['cam_param_pix_inv'].copy() | |
if k[0, 2] > 0: | |
k = np.linalg.inv(k) | |
ray = np.dot(k, np.array([sy-mesh.graph['woffset'], sx-mesh.graph['hoffset'], 1]).reshape(3, 1)) | |
point_3d = ray * np.abs(z) | |
point_3d = point_3d.flatten() | |
return point_3d | |
def generate_init_node(mesh, config, min_node_in_cc): | |
mesh_nodes = mesh.nodes | |
info_on_pix = {} | |
ccs = sorted(netx.connected_components(mesh), key = len, reverse=True) | |
remove_nodes = [] | |
for cc in ccs: | |
remove_flag = True if len(cc) < min_node_in_cc else False | |
if remove_flag is False: | |
for (nx, ny, nd) in cc: | |
info_on_pix[(nx, ny)] = [{'depth':nd, | |
'color':mesh_nodes[(nx, ny, nd)]['color'], | |
'synthesis':False, | |
'disp':mesh_nodes[(nx, ny, nd)]['disp']}] | |
else: | |
[remove_nodes.append((nx, ny, nd)) for (nx, ny, nd) in cc] | |
for node in remove_nodes: | |
far_nodes = [] if mesh_nodes[node].get('far') is None else mesh_nodes[node]['far'] | |
for far_node in far_nodes: | |
if mesh.has_node(far_node) and mesh_nodes[far_node].get('near') is not None and node in mesh_nodes[far_node]['near']: | |
mesh_nodes[far_node]['near'].remove(node) | |
near_nodes = [] if mesh_nodes[node].get('near') is None else mesh_nodes[node]['near'] | |
for near_node in near_nodes: | |
if mesh.has_node(near_node) and mesh_nodes[near_node].get('far') is not None and node in mesh_nodes[near_node]['far']: | |
mesh_nodes[near_node]['far'].remove(node) | |
[mesh.remove_node(node) for node in remove_nodes] | |
return mesh, info_on_pix | |
def get_neighbors(mesh, node): | |
return [*mesh.neighbors(node)] | |
def generate_face(mesh, info_on_pix, config): | |
H, W = mesh.graph['H'], mesh.graph['W'] | |
str_faces = [] | |
num_node = len(mesh.nodes) | |
ply_flag = config.get('save_ply') | |
def out_fmt(input, cur_id_b, cur_id_self, cur_id_a, ply_flag): | |
if ply_flag is True: | |
input.append(' '.join(['3', cur_id_b, cur_id_self, cur_id_a]) + '\n') | |
else: | |
input.append([cur_id_b, cur_id_self, cur_id_a]) | |
mesh_nodes = mesh.nodes | |
for node in mesh_nodes: | |
cur_id_self = mesh_nodes[node]['cur_id'] | |
ne_nodes = get_neighbors(mesh, node) | |
four_dir_nes = {'up': [], 'left': [], | |
'down': [], 'right': []} | |
for ne_node in ne_nodes: | |
store_tuple = [ne_node, mesh_nodes[ne_node]['cur_id']] | |
if ne_node[0] == node[0]: | |
if ne_node[1] == ne_node[1] - 1: | |
four_dir_nes['left'].append(store_tuple) | |
else: | |
four_dir_nes['right'].append(store_tuple) | |
else: | |
if ne_node[0] == ne_node[0] - 1: | |
four_dir_nes['up'].append(store_tuple) | |
else: | |
four_dir_nes['down'].append(store_tuple) | |
for node_a, cur_id_a in four_dir_nes['up']: | |
for node_b, cur_id_b in four_dir_nes['right']: | |
out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag) | |
for node_a, cur_id_a in four_dir_nes['right']: | |
for node_b, cur_id_b in four_dir_nes['down']: | |
out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag) | |
for node_a, cur_id_a in four_dir_nes['down']: | |
for node_b, cur_id_b in four_dir_nes['left']: | |
out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag) | |
for node_a, cur_id_a in four_dir_nes['left']: | |
for node_b, cur_id_b in four_dir_nes['up']: | |
out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag) | |
return str_faces | |
def reassign_floating_island(mesh, info_on_pix, image, depth): | |
H, W = mesh.graph['H'], mesh.graph['W'], | |
mesh_nodes = mesh.nodes | |
bord_up, bord_down = mesh.graph['bord_up'], mesh.graph['bord_down'] | |
bord_left, bord_right = mesh.graph['bord_left'], mesh.graph['bord_right'] | |
W = mesh.graph['W'] | |
lost_map = np.zeros((H, W)) | |
''' | |
(5) is_inside(x, y, xmin, xmax, ymin, ymax) : Check if a pixel(x, y) is inside the border. | |
(6) get_cross_nes(x, y) : Get the four cross neighbors of pixel(x, y). | |
''' | |
key_exist = lambda d, k: k in d | |
is_inside = lambda x, y, xmin, xmax, ymin, ymax: xmin <= x < xmax and ymin <= y < ymax | |
get_cross_nes = lambda x, y: [(x + 1, y), (x - 1, y), (x, y - 1), (x, y + 1)] | |
''' | |
(A) Highlight the pixels on isolated floating island. | |
(B) Number those isolated floating islands with connected component analysis. | |
(C) For each isolated island: | |
(1) Find its longest surrounded depth edge. | |
(2) Propogate depth from that depth edge to the pixels on the isolated island. | |
(3) Build the connection between the depth edge and that isolated island. | |
''' | |
for x in range(H): | |
for y in range(W): | |
if is_inside(x, y, bord_up, bord_down, bord_left, bord_right) and not(key_exist(info_on_pix, (x, y))): | |
lost_map[x, y] = 1 | |
_, label_lost_map = cv2.connectedComponents(lost_map.astype(np.uint8), connectivity=4) | |
mask = np.zeros((H, W)) | |
mask[bord_up:bord_down, bord_left:bord_right] = 1 | |
label_lost_map = (label_lost_map * mask).astype(np.int) | |
for i in range(1, label_lost_map.max()+1): | |
lost_xs, lost_ys = np.where(label_lost_map == i) | |
surr_edge_ids = {} | |
for lost_x, lost_y in zip(lost_xs, lost_ys): | |
if (lost_x, lost_y) == (295, 389) or (lost_x, lost_y) == (296, 389): | |
import pdb; pdb.set_trace() | |
for ne in get_cross_nes(lost_x, lost_y): | |
if key_exist(info_on_pix, ne): | |
for info in info_on_pix[ne]: | |
ne_node = (ne[0], ne[1], info['depth']) | |
if key_exist(mesh_nodes[ne_node], 'edge_id'): | |
edge_id = mesh_nodes[ne_node]['edge_id'] | |
surr_edge_ids[edge_id] = surr_edge_ids[edge_id] + [ne_node] if \ | |
key_exist(surr_edge_ids, edge_id) else [ne_node] | |
if len(surr_edge_ids) == 0: | |
continue | |
edge_id, edge_nodes = sorted([*surr_edge_ids.items()], key=lambda x: len(x[1]), reverse=True)[0] | |
edge_depth_map = np.zeros((H, W)) | |
for node in edge_nodes: | |
edge_depth_map[node[0], node[1]] = node[2] | |
lost_xs, lost_ys = np.where(label_lost_map == i) | |
while lost_xs.shape[0] > 0: | |
lost_xs, lost_ys = np.where(label_lost_map == i) | |
for lost_x, lost_y in zip(lost_xs, lost_ys): | |
propagated_depth = [] | |
real_nes = [] | |
for ne in get_cross_nes(lost_x, lost_y): | |
if not(is_inside(ne[0], ne[1], bord_up, bord_down, bord_left, bord_right)) or \ | |
edge_depth_map[ne[0], ne[1]] == 0: | |
continue | |
propagated_depth.append(edge_depth_map[ne[0], ne[1]]) | |
real_nes.append(ne) | |
if len(real_nes) == 0: | |
continue | |
reassign_depth = np.mean(propagated_depth) | |
label_lost_map[lost_x, lost_y] = 0 | |
edge_depth_map[lost_x, lost_y] = reassign_depth | |
depth[lost_x, lost_y] = -reassign_depth | |
mesh.add_node((lost_x, lost_y, reassign_depth), color=image[lost_x, lost_y], | |
synthesis=False, | |
disp=1./reassign_depth, | |
cc_id=set()) | |
info_on_pix[(lost_x, lost_y)] = [{'depth':reassign_depth, | |
'color':image[lost_x, lost_y], | |
'synthesis':False, | |
'disp':1./reassign_depth}] | |
new_connections = [((lost_x, lost_y, reassign_depth), | |
(ne[0], ne[1], edge_depth_map[ne[0], ne[1]])) for ne in real_nes] | |
mesh.add_edges_from(new_connections) | |
return mesh, info_on_pix, depth | |
def remove_node_feat(mesh, *feats): | |
mesh_nodes = mesh.nodes | |
for node in mesh_nodes: | |
for feat in feats: | |
mesh_nodes[node][feat] = None | |
return mesh | |
def update_status(mesh, info_on_pix, depth=None): | |
''' | |
(2) clear_node_feat(G, *fts) : Clear all the node feature on graph G. | |
(6) get_cross_nes(x, y) : Get the four cross neighbors of pixel(x, y). | |
''' | |
key_exist = lambda d, k: d.get(k) is not None | |
is_inside = lambda x, y, xmin, xmax, ymin, ymax: xmin <= x < xmax and ymin <= y < ymax | |
get_cross_nes = lambda x, y: [(x + 1, y), (x - 1, y), (x, y - 1), (x, y + 1)] | |
append_element = lambda d, k, x: d[k] + [x] if key_exist(d, k) else [x] | |
def clear_node_feat(G, fts): | |
le_nodes = G.nodes | |
for k in le_nodes: | |
v = le_nodes[k] | |
for ft in fts: | |
if ft in v: | |
v[ft] = None | |
clear_node_feat(mesh, ['edge_id', 'far', 'near']) | |
bord_up, bord_down = mesh.graph['bord_up'], mesh.graph['bord_down'] | |
bord_left, bord_right = mesh.graph['bord_left'], mesh.graph['bord_right'] | |
le_nodes = mesh.nodes | |
for node_key in le_nodes: | |
if mesh.neighbors(node_key).__length_hint__() == 4: | |
continue | |
four_nes = [xx for xx in get_cross_nes(node_key[0], node_key[1]) if | |
is_inside(xx[0], xx[1], bord_up, bord_down, bord_left, bord_right) and | |
xx in info_on_pix] | |
[four_nes.remove((ne_node[0], ne_node[1])) for ne_node in mesh.neighbors(node_key)] | |
for ne in four_nes: | |
for info in info_on_pix[ne]: | |
assert mesh.has_node((ne[0], ne[1], info['depth'])), "No node_key" | |
ind_node = le_nodes[node_key] | |
if abs(node_key[2]) > abs(info['depth']): | |
ind_node['near'] = append_element(ind_node, 'near', (ne[0], ne[1], info['depth'])) | |
else: | |
ind_node['far'] = append_element(ind_node, 'far', (ne[0], ne[1], info['depth'])) | |
if depth is not None: | |
for key, value in info_on_pix.items(): | |
if depth[key[0], key[1]] != abs(value[0]['depth']): | |
value[0]['disp'] = 1. / value[0]['depth'] | |
depth[key[0], key[1]] = abs(value[0]['depth']) | |
return mesh, depth, info_on_pix | |
else: | |
return mesh | |
def group_edges(LDI, config, image, remove_conflict_ordinal, spdb=False): | |
''' | |
(1) add_new_node(G, node) : add "node" to graph "G" | |
(2) add_new_edge(G, node_a, node_b) : add edge "node_a--node_b" to graph "G" | |
(3) exceed_thre(x, y, thre) : Check if difference between "x" and "y" exceed threshold "thre" | |
(4) key_exist(d, k) : Check if key "k' exists in dictionary "d" | |
(5) comm_opp_bg(G, x, y) : Check if node "x" and "y" in graph "G" treat the same opposite node as background | |
(6) comm_opp_fg(G, x, y) : Check if node "x" and "y" in graph "G" treat the same opposite node as foreground | |
''' | |
add_new_node = lambda G, node: None if G.has_node(node) else G.add_node(node) | |
add_new_edge = lambda G, node_a, node_b: None if G.has_edge(node_a, node_b) else G.add_edge(node_a, node_b) | |
exceed_thre = lambda x, y, thre: (abs(x) - abs(y)) > thre | |
key_exist = lambda d, k: d.get(k) is not None | |
comm_opp_bg = lambda G, x, y: key_exist(G.nodes[x], 'far') and key_exist(G.nodes[y], 'far') and \ | |
not(set(G.nodes[x]['far']).isdisjoint(set(G.nodes[y]['far']))) | |
comm_opp_fg = lambda G, x, y: key_exist(G.nodes[x], 'near') and key_exist(G.nodes[y], 'near') and \ | |
not(set(G.nodes[x]['near']).isdisjoint(set(G.nodes[y]['near']))) | |
discont_graph = netx.Graph() | |
''' | |
(A) Skip the pixel at image boundary, we don't want to deal with them. | |
(B) Identify discontinuity by the number of its neighbor(degree). | |
If the degree < 4(up/right/buttom/left). We will go through following steps: | |
(1) Add the discontinuity pixel "node" to graph "discont_graph". | |
(2) Find "node"'s cross neighbor(up/right/buttom/left) "ne_node". | |
- If the cross neighbor "ne_node" is a discontinuity pixel(degree("ne_node") < 4), | |
(a) add it to graph "discont_graph" and build the connection between "ne_node" and "node". | |
(b) label its cross neighbor as invalid pixels "inval_diag_candi" to avoid building | |
connection between original discontinuity pixel "node" and "inval_diag_candi". | |
- Otherwise, find "ne_node"'s cross neighbors, called diagonal candidate "diag_candi". | |
- The "diag_candi" is diagonal to the original discontinuity pixel "node". | |
- If "diag_candi" exists, go to step(3). | |
(3) A diagonal candidate "diag_candi" will be : | |
- added to the "discont_graph" if its degree < 4. | |
- connected to the original discontinuity pixel "node" if it satisfied either | |
one of following criterion: | |
(a) the difference of disparity between "diag_candi" and "node" is smaller than default threshold. | |
(b) the "diag_candi" and "node" face the same opposite pixel. (See. function "tear_edges") | |
(c) Both of "diag_candi" and "node" must_connect to each other. (See. function "combine_end_node") | |
(C) Aggregate each connected part in "discont_graph" into "discont_ccs" (A.K.A. depth edge). | |
''' | |
for node in LDI.nodes: | |
if not(LDI.graph['bord_up'] + 1 <= node[0] <= LDI.graph['bord_down'] - 2 and \ | |
LDI.graph['bord_left'] + 1 <= node[1] <= LDI.graph['bord_right'] - 2): | |
continue | |
neighbors = [*LDI.neighbors(node)] | |
if len(neighbors) < 4: | |
add_new_node(discont_graph, node) | |
diag_candi_anc, inval_diag_candi, discont_nes = set(), set(), set() | |
for ne_node in neighbors: | |
if len([*LDI.neighbors(ne_node)]) < 4: | |
add_new_node(discont_graph, ne_node) | |
add_new_edge(discont_graph, ne_node, node) | |
discont_nes.add(ne_node) | |
else: | |
diag_candi_anc.add(ne_node) | |
inval_diag_candi = set([inval_diagonal for ne_node in discont_nes for inval_diagonal in LDI.neighbors(ne_node) if \ | |
abs(inval_diagonal[0] - node[0]) < 2 and abs(inval_diagonal[1] - node[1]) < 2]) | |
for ne_node in diag_candi_anc: | |
if ne_node[0] == node[0]: | |
diagonal_xys = [[ne_node[0] + 1, ne_node[1]], [ne_node[0] - 1, ne_node[1]]] | |
elif ne_node[1] == node[1]: | |
diagonal_xys = [[ne_node[0], ne_node[1] + 1], [ne_node[0], ne_node[1] - 1]] | |
for diag_candi in LDI.neighbors(ne_node): | |
if [diag_candi[0], diag_candi[1]] in diagonal_xys and LDI.degree(diag_candi) < 4: | |
if diag_candi not in inval_diag_candi: | |
if not exceed_thre(1./node[2], 1./diag_candi[2], config['depth_threshold']) or \ | |
(comm_opp_bg(LDI, diag_candi, node) and comm_opp_fg(LDI, diag_candi, node)): | |
add_new_node(discont_graph, diag_candi) | |
add_new_edge(discont_graph, diag_candi, node) | |
if key_exist(LDI.nodes[diag_candi], 'must_connect') and node in LDI.nodes[diag_candi]['must_connect'] and \ | |
key_exist(LDI.nodes[node], 'must_connect') and diag_candi in LDI.nodes[node]['must_connect']: | |
add_new_node(discont_graph, diag_candi) | |
add_new_edge(discont_graph, diag_candi, node) | |
if spdb == True: | |
import pdb; pdb.set_trace() | |
discont_ccs = [*netx.connected_components(discont_graph)] | |
''' | |
In some corner case, a depth edge "discont_cc" will contain both | |
foreground(FG) and background(BG) pixels. This violate the assumption that | |
a depth edge can only composite by one type of pixel(FG or BG). | |
We need to further divide this depth edge into several sub-part so that the | |
assumption is satisfied. | |
(A) A depth edge is invalid if both of its "far_flag"(BG) and | |
"near_flag"(FG) are True. | |
(B) If the depth edge is invalid, we need to do: | |
(1) Find the role("oridinal") of each pixel on the depth edge. | |
"-1" --> Its opposite pixels has smaller depth(near) than it. | |
It is a backgorund pixel. | |
"+1" --> Its opposite pixels has larger depth(far) than it. | |
It is a foregorund pixel. | |
"0" --> Some of opposite pixels has larger depth(far) than it, | |
and some has smaller pixel than it. | |
It is an ambiguous pixel. | |
(2) For each pixel "discont_node", check if its neigbhors' roles are consistent. | |
- If not, break the connection between the neighbor "ne_node" that has a role | |
different from "discont_node". | |
- If yes, remove all the role that are inconsistent to its neighbors "ne_node". | |
(3) Connected component analysis to re-identified those divided depth edge. | |
(C) Aggregate each connected part in "discont_graph" into "discont_ccs" (A.K.A. depth edge). | |
''' | |
if remove_conflict_ordinal: | |
new_discont_ccs = [] | |
num_new_cc = 0 | |
for edge_id, discont_cc in enumerate(discont_ccs): | |
near_flag = False | |
far_flag = False | |
for discont_node in discont_cc: | |
near_flag = True if key_exist(LDI.nodes[discont_node], 'far') else near_flag | |
far_flag = True if key_exist(LDI.nodes[discont_node], 'near') else far_flag | |
if far_flag and near_flag: | |
break | |
if far_flag and near_flag: | |
for discont_node in discont_cc: | |
discont_graph.nodes[discont_node]['ordinal'] = \ | |
np.array([key_exist(LDI.nodes[discont_node], 'far'), | |
key_exist(LDI.nodes[discont_node], 'near')]) * \ | |
np.array([-1, 1]) | |
discont_graph.nodes[discont_node]['ordinal'] = \ | |
np.sum(discont_graph.nodes[discont_node]['ordinal']) | |
remove_nodes, remove_edges = [], [] | |
for discont_node in discont_cc: | |
ordinal_relation = np.sum([discont_graph.nodes[xx]['ordinal'] \ | |
for xx in discont_graph.neighbors(discont_node)]) | |
near_side = discont_graph.nodes[discont_node]['ordinal'] <= 0 | |
if abs(ordinal_relation) < len([*discont_graph.neighbors(discont_node)]): | |
remove_nodes.append(discont_node) | |
for ne_node in discont_graph.neighbors(discont_node): | |
remove_flag = (near_side and not(key_exist(LDI.nodes[ne_node], 'far'))) or \ | |
(not near_side and not(key_exist(LDI.nodes[ne_node], 'near'))) | |
remove_edges += [(discont_node, ne_node)] if remove_flag else [] | |
else: | |
if near_side and key_exist(LDI.nodes[discont_node], 'near'): | |
LDI.nodes[discont_node].pop('near') | |
elif not(near_side) and key_exist(LDI.nodes[discont_node], 'far'): | |
LDI.nodes[discont_node].pop('far') | |
discont_graph.remove_edges_from(remove_edges) | |
sub_mesh = discont_graph.subgraph(list(discont_cc)).copy() | |
sub_discont_ccs = [*netx.connected_components(sub_mesh)] | |
is_redun_near = lambda xx: len(xx) == 1 and xx[0] in remove_nodes and key_exist(LDI.nodes[xx[0]], 'far') | |
for sub_discont_cc in sub_discont_ccs: | |
if is_redun_near(list(sub_discont_cc)): | |
LDI.nodes[list(sub_discont_cc)[0]].pop('far') | |
new_discont_ccs.append(sub_discont_cc) | |
else: | |
new_discont_ccs.append(discont_cc) | |
discont_ccs = new_discont_ccs | |
new_discont_ccs = None | |
if spdb == True: | |
import pdb; pdb.set_trace() | |
for edge_id, edge_cc in enumerate(discont_ccs): | |
for node in edge_cc: | |
LDI.nodes[node]['edge_id'] = edge_id | |
return discont_ccs, LDI, discont_graph | |
def combine_end_node(mesh, edge_mesh, edge_ccs, depth): | |
import collections | |
mesh_nodes = mesh.nodes | |
connect_dict = dict() | |
for valid_edge_id, valid_edge_cc in enumerate(edge_ccs): | |
connect_info = [] | |
for valid_edge_node in valid_edge_cc: | |
single_connect = set() | |
for ne_node in mesh.neighbors(valid_edge_node): | |
if mesh_nodes[ne_node].get('far') is not None: | |
for fn in mesh_nodes[ne_node].get('far'): | |
if mesh.has_node(fn) and mesh_nodes[fn].get('edge_id') is not None: | |
single_connect.add(mesh_nodes[fn]['edge_id']) | |
if mesh_nodes[ne_node].get('near') is not None: | |
for fn in mesh_nodes[ne_node].get('near'): | |
if mesh.has_node(fn) and mesh_nodes[fn].get('edge_id') is not None: | |
single_connect.add(mesh_nodes[fn]['edge_id']) | |
connect_info.extend([*single_connect]) | |
connect_dict[valid_edge_id] = collections.Counter(connect_info) | |
end_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
edge_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1 | |
for valid_edge_id, valid_edge_cc in enumerate(edge_ccs): | |
for valid_edge_node in valid_edge_cc: | |
edge_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id | |
if len([*edge_mesh.neighbors(valid_edge_node)]) == 1: | |
num_ne = 1 | |
if num_ne == 1: | |
end_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_node[2] | |
nxs, nys = np.where(end_maps != 0) | |
invalid_nodes = set() | |
for nx, ny in zip(nxs, nys): | |
if mesh.has_node((nx, ny, end_maps[nx, ny])) is False: | |
invalid_nodes.add((nx, ny)) | |
continue | |
four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \ | |
if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \ | |
end_maps[xx[0], xx[1]] != 0] | |
mesh_nes = [*mesh.neighbors((nx, ny, end_maps[nx, ny]))] | |
remove_num = 0 | |
for fne in four_nes: | |
if (fne[0], fne[1], end_maps[fne[0], fne[1]]) in mesh_nes: | |
remove_num += 1 | |
if remove_num == len(four_nes): | |
invalid_nodes.add((nx, ny)) | |
for invalid_node in invalid_nodes: | |
end_maps[invalid_node[0], invalid_node[1]] = 0 | |
nxs, nys = np.where(end_maps != 0) | |
invalid_nodes = set() | |
for nx, ny in zip(nxs, nys): | |
if mesh_nodes[(nx, ny, end_maps[nx, ny])].get('edge_id') is None: | |
continue | |
else: | |
self_id = mesh_nodes[(nx, ny, end_maps[nx, ny])].get('edge_id') | |
self_connect = connect_dict[self_id] if connect_dict.get(self_id) is not None else dict() | |
four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \ | |
if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \ | |
end_maps[xx[0], xx[1]] != 0] | |
for fne in four_nes: | |
if mesh_nodes[(fne[0], fne[1], end_maps[fne[0], fne[1]])].get('edge_id') is None: | |
continue | |
else: | |
ne_id = mesh_nodes[(fne[0], fne[1], end_maps[fne[0], fne[1]])]['edge_id'] | |
if self_connect.get(ne_id) is None or self_connect.get(ne_id) == 1: | |
continue | |
else: | |
invalid_nodes.add((nx, ny)) | |
for invalid_node in invalid_nodes: | |
end_maps[invalid_node[0], invalid_node[1]] = 0 | |
nxs, nys = np.where(end_maps != 0) | |
invalid_nodes = set() | |
for nx, ny in zip(nxs, nys): | |
four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \ | |
if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \ | |
end_maps[xx[0], xx[1]] != 0] | |
for fne in four_nes: | |
if mesh.has_node((fne[0], fne[1], end_maps[fne[0], fne[1]])): | |
node_a, node_b = (fne[0], fne[1], end_maps[fne[0], fne[1]]), (nx, ny, end_maps[nx, ny]) | |
mesh.add_edge(node_a, node_b) | |
mesh_nodes[node_b]['must_connect'] = set() if mesh_nodes[node_b].get('must_connect') is None else mesh_nodes[node_b]['must_connect'] | |
mesh_nodes[node_b]['must_connect'].add(node_a) | |
mesh_nodes[node_b]['must_connect'] |= set([xx for xx in [*edge_mesh.neighbors(node_a)] if \ | |
(xx[0] - node_b[0]) < 2 and (xx[1] - node_b[1]) < 2]) | |
mesh_nodes[node_a]['must_connect'] = set() if mesh_nodes[node_a].get('must_connect') is None else mesh_nodes[node_a]['must_connect'] | |
mesh_nodes[node_a]['must_connect'].add(node_b) | |
mesh_nodes[node_a]['must_connect'] |= set([xx for xx in [*edge_mesh.neighbors(node_b)] if \ | |
(xx[0] - node_a[0]) < 2 and (xx[1] - node_a[1]) < 2]) | |
invalid_nodes.add((nx, ny)) | |
for invalid_node in invalid_nodes: | |
end_maps[invalid_node[0], invalid_node[1]] = 0 | |
return mesh | |
def remove_redundant_edge(mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=1000, invalid=False, spdb=False): | |
point_to_amount = {} | |
point_to_id = {} | |
end_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1 | |
for valid_edge_id, valid_edge_cc in enumerate(edge_ccs): | |
for valid_edge_node in valid_edge_cc: | |
point_to_amount[valid_edge_node] = len(valid_edge_cc) | |
point_to_id[valid_edge_node] = valid_edge_id | |
if edge_mesh.has_node(valid_edge_node) is True: | |
if len([*edge_mesh.neighbors(valid_edge_node)]) == 1: | |
end_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id | |
nxs, nys = np.where(end_maps > -1) | |
point_to_adjoint = {} | |
for nx, ny in zip(nxs, nys): | |
adjoint_edges = set([end_maps[x, y] for x, y in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)] if end_maps[x, y] != -1]) | |
point_to_adjoint[end_maps[nx, ny]] = (point_to_adjoint[end_maps[nx, ny]] | adjoint_edges) if point_to_adjoint.get(end_maps[nx, ny]) is not None else adjoint_edges | |
valid_edge_ccs = filter_edge(mesh, edge_ccs, config, invalid=invalid) | |
edge_canvas = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1 | |
for valid_edge_id, valid_edge_cc in enumerate(valid_edge_ccs): | |
for valid_edge_node in valid_edge_cc: | |
edge_canvas[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id | |
if spdb is True: | |
plt.imshow(edge_canvas); plt.show() | |
import pdb; pdb.set_trace() | |
for valid_edge_id, valid_edge_cc in enumerate(valid_edge_ccs): | |
end_number = 0 | |
four_end_number = 0 | |
eight_end_number = 0 | |
db_eight_end_number = 0 | |
if len(valid_edge_cc) > redundant_number: | |
continue | |
for valid_edge_node in valid_edge_cc: | |
if len([*edge_mesh.neighbors(valid_edge_node)]) == 3: | |
break | |
elif len([*edge_mesh.neighbors(valid_edge_node)]) == 1: | |
hx, hy, hz = valid_edge_node | |
if invalid is False: | |
eight_nes = [(x, y) for x, y 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 info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id] | |
if len(eight_nes) == 0: | |
end_number += 1 | |
if invalid is True: | |
four_nes = []; eight_nes = []; db_eight_nes = [] | |
four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \ | |
if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id] | |
eight_nes = [(x, y) for x, y 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 info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id] | |
db_eight_nes = [(x, y) for x in range(hx - 2, hx + 3) for y in range(hy - 2, hy + 3) \ | |
if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id and (x, y) != (hx, hy)] | |
if len(four_nes) == 0 or len(eight_nes) == 0: | |
end_number += 1 | |
if len(four_nes) == 0: | |
four_end_number += 1 | |
if len(eight_nes) == 0: | |
eight_end_number += 1 | |
if len(db_eight_nes) == 0: | |
db_eight_end_number += 1 | |
elif len([*edge_mesh.neighbors(valid_edge_node)]) == 0: | |
hx, hy, hz = valid_edge_node | |
four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \ | |
if info_on_pix.get((x, y)) is not None and \ | |
mesh.has_edge(valid_edge_node, (x, y, info_on_pix[(x, y)][0]['depth'])) is False] | |
for ne in four_nes: | |
try: | |
if invalid is True or (point_to_amount.get(ne) is None or point_to_amount[ne] < redundant_number) or \ | |
point_to_id[ne] in point_to_adjoint.get(point_to_id[valid_edge_node], set()): | |
mesh.add_edge(valid_edge_node, ne) | |
except: | |
import pdb; pdb.set_trace() | |
if (invalid is not True and end_number >= 1) or (invalid is True and end_number >= 2 and eight_end_number >= 1 and db_eight_end_number >= 1): | |
for valid_edge_node in valid_edge_cc: | |
hx, hy, _ = valid_edge_node | |
four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \ | |
if info_on_pix.get((x, y)) is not None and \ | |
mesh.has_edge(valid_edge_node, (x, y, info_on_pix[(x, y)][0]['depth'])) is False and \ | |
(edge_canvas[x, y] == -1 or edge_canvas[x, y] == valid_edge_id)] | |
for ne in four_nes: | |
if invalid is True or (point_to_amount.get(ne) is None or point_to_amount[ne] < redundant_number) or \ | |
point_to_id[ne] in point_to_adjoint.get(point_to_id[valid_edge_node], set()): | |
mesh.add_edge(valid_edge_node, ne) | |
return mesh | |
def judge_dangle(mark, mesh, node): | |
if not (1 <= node[0] < mesh.graph['H']-1) or not(1 <= node[1] < mesh.graph['W']-1): | |
return mark | |
mesh_neighbors = [*mesh.neighbors(node)] | |
mesh_neighbors = [xx for xx in mesh_neighbors if 0 < xx[0] < mesh.graph['H'] - 1 and 0 < xx[1] < mesh.graph['W'] - 1] | |
if len(mesh_neighbors) >= 3: | |
return mark | |
elif len(mesh_neighbors) <= 1: | |
mark[node[0], node[1]] = (len(mesh_neighbors) + 1) | |
else: | |
dan_ne_node_a = mesh_neighbors[0] | |
dan_ne_node_b = mesh_neighbors[1] | |
if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \ | |
abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1: | |
mark[node[0], node[1]] = 3 | |
return mark | |
def remove_dangling(mesh, edge_ccs, edge_mesh, info_on_pix, image, depth, config): | |
tmp_edge_ccs = copy.deepcopy(edge_ccs) | |
for edge_cc_id, valid_edge_cc in enumerate(tmp_edge_ccs): | |
if len(valid_edge_cc) > 1 or len(valid_edge_cc) == 0: | |
continue | |
single_edge_node = [*valid_edge_cc][0] | |
hx, hy, hz = single_edge_node | |
eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y 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 info_on_pix.get((x, y)) is not None]) | |
four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \ | |
if info_on_pix.get((x, y)) is not None] | |
sub_mesh = mesh.subgraph(eight_nes).copy() | |
ccs = netx.connected_components(sub_mesh) | |
four_ccs = [] | |
for cc_id, _cc in enumerate(ccs): | |
four_ccs.append(set()) | |
for cc_node in _cc: | |
if abs(cc_node[0] - hx) + abs(cc_node[1] - hy) < 2: | |
four_ccs[cc_id].add(cc_node) | |
largest_cc = sorted(four_ccs, key=lambda x: (len(x), -np.sum([abs(xx[2] - hz) for xx in x])))[-1] | |
if len(largest_cc) < 2: | |
for ne in four_nes: | |
mesh.add_edge(single_edge_node, ne) | |
else: | |
mesh.remove_edges_from([(single_edge_node, ne) for ne in mesh.neighbors(single_edge_node)]) | |
new_depth = np.mean([xx[2] for xx in largest_cc]) | |
info_on_pix[(hx, hy)][0]['depth'] = new_depth | |
info_on_pix[(hx, hy)][0]['disp'] = 1./new_depth | |
new_node = (hx, hy, new_depth) | |
mesh = refresh_node(single_edge_node, mesh.node[single_edge_node], new_node, dict(), mesh) | |
edge_ccs[edge_cc_id] = set([new_node]) | |
for ne in largest_cc: | |
mesh.add_edge(new_node, ne) | |
mark = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
for edge_idx, edge_cc in enumerate(edge_ccs): | |
for edge_node in edge_cc: | |
if not (mesh.graph['bord_up'] <= edge_node[0] < mesh.graph['bord_down']-1) or \ | |
not (mesh.graph['bord_left'] <= edge_node[1] < mesh.graph['bord_right']-1): | |
continue | |
mesh_neighbors = [*mesh.neighbors(edge_node)] | |
mesh_neighbors = [xx for xx in mesh_neighbors \ | |
if mesh.graph['bord_up'] < xx[0] < mesh.graph['bord_down'] - 1 and \ | |
mesh.graph['bord_left'] < xx[1] < mesh.graph['bord_right'] - 1] | |
if len([*mesh.neighbors(edge_node)]) >= 3: | |
continue | |
elif len([*mesh.neighbors(edge_node)]) <= 1: | |
mark[edge_node[0], edge_node[1]] += (len([*mesh.neighbors(edge_node)]) + 1) | |
else: | |
dan_ne_node_a = [*mesh.neighbors(edge_node)][0] | |
dan_ne_node_b = [*mesh.neighbors(edge_node)][1] | |
if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \ | |
abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1: | |
mark[edge_node[0], edge_node[1]] += 3 | |
mxs, mys = np.where(mark == 1) | |
conn_0_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \ | |
if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']))] | |
mxs, mys = np.where(mark == 2) | |
conn_1_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \ | |
if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']))] | |
for node in conn_0_nodes: | |
hx, hy = node[0], node[1] | |
four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \ | |
if info_on_pix.get((x, y)) is not None] | |
re_depth = {'value' : 0, 'count': 0} | |
for ne in four_nes: | |
mesh.add_edge(node, ne) | |
re_depth['value'] += cc_node[2] | |
re_depth['count'] += 1. | |
re_depth = re_depth['value'] / re_depth['count'] | |
mapping_dict = {node: (node[0], node[1], re_depth)} | |
info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh) | |
depth[node[0], node[1]] = abs(re_depth) | |
mark[node[0], node[1]] = 0 | |
for node in conn_1_nodes: | |
hx, hy = node[0], node[1] | |
eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y 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 info_on_pix.get((x, y)) is not None]) | |
self_nes = set([ne2 for ne1 in mesh.neighbors(node) for ne2 in mesh.neighbors(ne1) if ne2 in eight_nes]) | |
eight_nes = [*(eight_nes - self_nes)] | |
sub_mesh = mesh.subgraph(eight_nes).copy() | |
ccs = netx.connected_components(sub_mesh) | |
largest_cc = sorted(ccs, key=lambda x: (len(x), -np.sum([abs(xx[0] - node[0]) + abs(xx[1] - node[1]) for xx in x])))[-1] | |
mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)]) | |
re_depth = {'value' : 0, 'count': 0} | |
for cc_node in largest_cc: | |
if cc_node[0] == node[0] and cc_node[1] == node[1]: | |
continue | |
re_depth['value'] += cc_node[2] | |
re_depth['count'] += 1. | |
if abs(cc_node[0] - node[0]) + abs(cc_node[1] - node[1]) < 2: | |
mesh.add_edge(cc_node, node) | |
try: | |
re_depth = re_depth['value'] / re_depth['count'] | |
except: | |
re_depth = node[2] | |
renode = (node[0], node[1], re_depth) | |
mapping_dict = {node: renode} | |
info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh) | |
depth[node[0], node[1]] = abs(re_depth) | |
mark[node[0], node[1]] = 0 | |
edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, renode, mark) | |
mxs, mys = np.where(mark == 3) | |
conn_2_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \ | |
if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth'])) and \ | |
mesh.degree((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth'])) == 2] | |
sub_mesh = mesh.subgraph(conn_2_nodes).copy() | |
ccs = netx.connected_components(sub_mesh) | |
for cc in ccs: | |
candidate_nodes = [xx for xx in cc if sub_mesh.degree(xx) == 1] | |
for node in candidate_nodes: | |
if mesh.has_node(node) is False: | |
continue | |
ne_node = [xx for xx in mesh.neighbors(node) if xx not in cc][0] | |
hx, hy = node[0], node[1] | |
eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y 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 info_on_pix.get((x, y)) is not None and (x, y, info_on_pix[(x, y)][0]['depth']) not in cc]) | |
ne_sub_mesh = mesh.subgraph(eight_nes).copy() | |
ne_ccs = netx.connected_components(ne_sub_mesh) | |
try: | |
ne_cc = [ne_cc for ne_cc in ne_ccs if ne_node in ne_cc][0] | |
except: | |
import pdb; pdb.set_trace() | |
largest_cc = [xx for xx in ne_cc if abs(xx[0] - node[0]) + abs(xx[1] - node[1]) == 1] | |
mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)]) | |
re_depth = {'value' : 0, 'count': 0} | |
for cc_node in largest_cc: | |
re_depth['value'] += cc_node[2] | |
re_depth['count'] += 1. | |
mesh.add_edge(cc_node, node) | |
try: | |
re_depth = re_depth['value'] / re_depth['count'] | |
except: | |
re_depth = node[2] | |
renode = (node[0], node[1], re_depth) | |
mapping_dict = {node: renode} | |
info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh) | |
depth[node[0], node[1]] = abs(re_depth) | |
mark[node[0], node[1]] = 0 | |
edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, renode, mark) | |
break | |
if len(cc) == 1: | |
node = [node for node in cc][0] | |
hx, hy = node[0], node[1] | |
nine_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx, hy), (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 info_on_pix.get((x, y)) is not None and mesh.has_node((x, y, info_on_pix[(x, y)][0]['depth']))]) | |
ne_sub_mesh = mesh.subgraph(nine_nes).copy() | |
ne_ccs = netx.connected_components(ne_sub_mesh) | |
for ne_cc in ne_ccs: | |
if node in ne_cc: | |
re_depth = {'value' : 0, 'count': 0} | |
for ne in ne_cc: | |
if abs(ne[0] - node[0]) + abs(ne[1] - node[1]) == 1: | |
mesh.add_edge(node, ne) | |
re_depth['value'] += ne[2] | |
re_depth['count'] += 1. | |
re_depth = re_depth['value'] / re_depth['count'] | |
mapping_dict = {node: (node[0], node[1], re_depth)} | |
info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh) | |
depth[node[0], node[1]] = abs(re_depth) | |
mark[node[0], node[1]] = 0 | |
return mesh, info_on_pix, edge_mesh, depth, mark | |
def context_and_holes(mesh, edge_ccs, config, specific_edge_id, specific_edge_loc, depth_feat_model, | |
connect_points_ccs=None, inpaint_iter=0, filter_edge=False, vis_edge_id=None): | |
edge_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1 | |
mask_info = {} | |
for edge_id, edge_cc in enumerate(edge_ccs): | |
for edge_node in edge_cc: | |
edge_maps[edge_node[0], edge_node[1]] = edge_id | |
context_ccs = [set() for x in range(len(edge_ccs))] | |
extend_context_ccs = [set() for x in range(len(edge_ccs))] | |
extend_erode_context_ccs = [set() for x in range(len(edge_ccs))] | |
extend_edge_ccs = [set() for x in range(len(edge_ccs))] | |
accomp_extend_context_ccs = [set() for x in range(len(edge_ccs))] | |
erode_context_ccs = [set() for x in range(len(edge_ccs))] | |
broken_mask_ccs = [set() for x in range(len(edge_ccs))] | |
invalid_extend_edge_ccs = [set() for x in range(len(edge_ccs))] | |
intouched_ccs = [set() for x in range(len(edge_ccs))] | |
redundant_ccs = [set() for x in range(len(edge_ccs))] | |
if inpaint_iter == 0: | |
background_thickness = config['background_thickness'] | |
context_thickness = config['context_thickness'] | |
else: | |
background_thickness = config['background_thickness_2'] | |
context_thickness = config['context_thickness_2'] | |
mesh_nodes = mesh.nodes | |
for edge_id, edge_cc in enumerate(edge_ccs): | |
if context_thickness == 0 or (len(specific_edge_id) > 0 and edge_id not in specific_edge_id): | |
continue | |
edge_group = {} | |
for edge_node in edge_cc: | |
far_nodes = mesh_nodes[edge_node].get('far') | |
if far_nodes is None: | |
continue | |
for far_node in far_nodes: | |
if far_node in edge_cc: | |
continue | |
context_ccs[edge_id].add(far_node) | |
if mesh_nodes[far_node].get('edge_id') is not None: | |
if edge_group.get(mesh_nodes[far_node]['edge_id']) is None: | |
edge_group[mesh_nodes[far_node]['edge_id']] = set() | |
edge_group[mesh_nodes[far_node]['edge_id']].add(far_node) | |
if len(edge_cc) > 2: | |
for edge_key in [*edge_group.keys()]: | |
if len(edge_group[edge_key]) == 1: | |
context_ccs[edge_id].remove([*edge_group[edge_key]][0]) | |
for edge_id, edge_cc in enumerate(edge_ccs): | |
if inpaint_iter != 0: | |
continue | |
tmp_intouched_nodes = set() | |
for edge_node in edge_cc: | |
raw_intouched_nodes = set(mesh_nodes[edge_node].get('near')) if mesh_nodes[edge_node].get('near') is not None else set() | |
tmp_intouched_nodes |= set([xx for xx in raw_intouched_nodes if mesh_nodes[xx].get('edge_id') is not None and \ | |
len(context_ccs[mesh_nodes[xx].get('edge_id')]) > 0]) | |
intouched_ccs[edge_id] |= tmp_intouched_nodes | |
tmp_intouched_nodes = None | |
mask_ccs = copy.deepcopy(edge_ccs) | |
forbidden_len = 3 | |
forbidden_map = np.ones((mesh.graph['H'] - forbidden_len, mesh.graph['W'] - forbidden_len)) | |
forbidden_map = np.pad(forbidden_map, ((forbidden_len, forbidden_len), (forbidden_len, forbidden_len)), mode='constant').astype(np.bool) | |
cur_tmp_mask_map = np.zeros_like(forbidden_map).astype(np.bool) | |
passive_background = 10 if 10 is not None else background_thickness | |
passive_context = 1 if 1 is not None else context_thickness | |
for edge_id, edge_cc in enumerate(edge_ccs): | |
cur_mask_cc = None; cur_mask_cc = [] | |
cur_context_cc = None; cur_context_cc = [] | |
cur_accomp_near_cc = None; cur_accomp_near_cc = [] | |
cur_invalid_extend_edge_cc = None; cur_invalid_extend_edge_cc = [] | |
cur_comp_far_cc = None; cur_comp_far_cc = [] | |
tmp_erode = [] | |
if len(context_ccs[edge_id]) == 0 or (len(specific_edge_id) > 0 and edge_id not in specific_edge_id): | |
continue | |
for i in range(max(background_thickness, context_thickness)): | |
cur_tmp_mask_map.fill(False) | |
if i == 0: | |
tmp_mask_nodes = copy.deepcopy(mask_ccs[edge_id]) | |
tmp_intersect_nodes = [] | |
tmp_intersect_context_nodes = [] | |
mask_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
context_depth = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
comp_cnt_depth = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
connect_map = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
for node in tmp_mask_nodes: | |
mask_map[node[0], node[1]] = True | |
depth_count = 0 | |
if mesh_nodes[node].get('far') is not None: | |
for comp_cnt_node in mesh_nodes[node]['far']: | |
comp_cnt_depth[node[0], node[1]] += abs(comp_cnt_node[2]) | |
depth_count += 1 | |
if depth_count > 0: | |
comp_cnt_depth[node[0], node[1]] = comp_cnt_depth[node[0], node[1]] / depth_count | |
connect_node = [] | |
if mesh_nodes[node].get('connect_point_id') is not None: | |
connect_node.append(mesh_nodes[node]['connect_point_id']) | |
connect_point_id = np.bincount(connect_node).argmax() if len(connect_node) > 0 else -1 | |
if connect_point_id > -1 and connect_points_ccs is not None: | |
for xx in connect_points_ccs[connect_point_id]: | |
if connect_map[xx[0], xx[1]] == 0: | |
connect_map[xx[0], xx[1]] = xx[2] | |
if mesh_nodes[node].get('connect_point_exception') is not None: | |
for xx in mesh_nodes[node]['connect_point_exception']: | |
if connect_map[xx[0], xx[1]] == 0: | |
connect_map[xx[0], xx[1]] = xx[2] | |
tmp_context_nodes = [*context_ccs[edge_id]] | |
tmp_erode.append([*context_ccs[edge_id]]) | |
context_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
if (context_map.astype(np.uint8) * mask_map.astype(np.uint8)).max() > 0: | |
import pdb; pdb.set_trace() | |
for node in tmp_context_nodes: | |
context_map[node[0], node[1]] = True | |
context_depth[node[0], node[1]] = node[2] | |
context_map[mask_map == True] = False | |
if (context_map.astype(np.uint8) * mask_map.astype(np.uint8)).max() > 0: | |
import pdb; pdb.set_trace() | |
tmp_intouched_nodes = [*intouched_ccs[edge_id]] | |
intouched_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
for node in tmp_intouched_nodes: intouched_map[node[0], node[1]] = True | |
intouched_map[mask_map == True] = False | |
tmp_redundant_nodes = set() | |
tmp_noncont_nodes = set() | |
noncont_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
intersect_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
intersect_context_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool) | |
if i > passive_background and inpaint_iter == 0: | |
new_tmp_intersect_nodes = None | |
new_tmp_intersect_nodes = [] | |
for node in tmp_intersect_nodes: | |
nes = mesh.neighbors(node) | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True and \ | |
bool(intouched_map[ne[0], ne[1]]) is False and\ | |
bool(intersect_map[ne[0], ne[1]]) is False and\ | |
bool(intersect_context_map[ne[0], ne[1]]) is False: | |
break_flag = False | |
if (i - passive_background) % 2 == 0 and (i - passive_background) % 8 != 0: | |
four_nes = [xx for xx in[[ne[0] - 1, ne[1]], [ne[0] + 1, ne[1]], [ne[0], ne[1] - 1], [ne[0], ne[1] + 1]] \ | |
if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W']] | |
for fne in four_nes: | |
if bool(mask_map[fne[0], fne[1]]) is True: | |
break_flag = True | |
break | |
if break_flag is True: | |
continue | |
intersect_map[ne[0], ne[1]] = True | |
new_tmp_intersect_nodes.append(ne) | |
tmp_intersect_nodes = None | |
tmp_intersect_nodes = new_tmp_intersect_nodes | |
if i > passive_context and inpaint_iter == 1: | |
new_tmp_intersect_context_nodes = None | |
new_tmp_intersect_context_nodes = [] | |
for node in tmp_intersect_context_nodes: | |
nes = mesh.neighbors(node) | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True and \ | |
bool(intouched_map[ne[0], ne[1]]) is False and\ | |
bool(intersect_map[ne[0], ne[1]]) is False and \ | |
bool(intersect_context_map[ne[0], ne[1]]) is False: | |
intersect_context_map[ne[0], ne[1]] = True | |
new_tmp_intersect_context_nodes.append(ne) | |
tmp_intersect_context_nodes = None | |
tmp_intersect_context_nodes = new_tmp_intersect_context_nodes | |
new_tmp_mask_nodes = None | |
new_tmp_mask_nodes = [] | |
for node in tmp_mask_nodes: | |
four_nes = {xx:[] for xx in [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)] if \ | |
0 <= xx[0] < connect_map.shape[0] and 0 <= xx[1] < connect_map.shape[1]} | |
if inpaint_iter > 0: | |
for ne in four_nes.keys(): | |
if connect_map[ne[0], ne[1]] == True: | |
tmp_context_nodes.append((ne[0], ne[1], connect_map[ne[0], ne[1]])) | |
context_map[ne[0], ne[1]] = True | |
nes = mesh.neighbors(node) | |
if inpaint_iter > 0: | |
for ne in nes: four_nes[(ne[0], ne[1])].append(ne[2]) | |
nes = [] | |
for kfne, vfnes in four_nes.items(): vfnes.sort(key = lambda xx: abs(xx), reverse=True) | |
for kfne, vfnes in four_nes.items(): | |
for vfne in vfnes: nes.append((kfne[0], kfne[1], vfne)) | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True and \ | |
bool(intouched_map[ne[0], ne[1]]) is False and \ | |
bool(intersect_map[ne[0], ne[1]]) is False and \ | |
bool(intersect_context_map[ne[0], ne[1]]) is False: | |
if i == passive_background and inpaint_iter == 0: | |
if np.any(context_map[max(ne[0] - 1, 0):min(ne[0] + 2, mesh.graph['H']), max(ne[1] - 1, 0):min(ne[1] + 2, mesh.graph['W'])]) == True: | |
intersect_map[ne[0], ne[1]] = True | |
tmp_intersect_nodes.append(ne) | |
continue | |
if i < background_thickness: | |
if inpaint_iter == 0: | |
cur_mask_cc.append(ne) | |
elif mesh_nodes[ne].get('inpaint_id') == 1: | |
cur_mask_cc.append(ne) | |
else: | |
continue | |
mask_ccs[edge_id].add(ne) | |
if inpaint_iter == 0: | |
if comp_cnt_depth[node[0], node[1]] > 0 and comp_cnt_depth[ne[0], ne[1]] == 0: | |
comp_cnt_depth[ne[0], ne[1]] = comp_cnt_depth[node[0], node[1]] | |
if mesh_nodes[ne].get('far') is not None: | |
for comp_far_node in mesh_nodes[ne]['far']: | |
cur_comp_far_cc.append(comp_far_node) | |
cur_accomp_near_cc.append(ne) | |
cur_invalid_extend_edge_cc.append(comp_far_node) | |
if mesh_nodes[ne].get('edge_id') is not None and \ | |
len(context_ccs[mesh_nodes[ne].get('edge_id')]) > 0: | |
intouched_fars = set(mesh_nodes[ne].get('far')) if mesh_nodes[ne].get('far') is not None else set() | |
accum_intouched_fars = set(intouched_fars) | |
for intouched_far in intouched_fars: | |
accum_intouched_fars |= set([*mesh.neighbors(intouched_far)]) | |
for intouched_far in accum_intouched_fars: | |
if bool(mask_map[intouched_far[0], intouched_far[1]]) is True or \ | |
bool(context_map[intouched_far[0], intouched_far[1]]) is True: | |
continue | |
tmp_redundant_nodes.add(intouched_far) | |
intouched_map[intouched_far[0], intouched_far[1]] = True | |
if mesh_nodes[ne].get('near') is not None: | |
intouched_nears = set(mesh_nodes[ne].get('near')) | |
for intouched_near in intouched_nears: | |
if bool(mask_map[intouched_near[0], intouched_near[1]]) is True or \ | |
bool(context_map[intouched_near[0], intouched_near[1]]) is True: | |
continue | |
tmp_redundant_nodes.add(intouched_near) | |
intouched_map[intouched_near[0], intouched_near[1]] = True | |
if not (mesh_nodes[ne].get('inpaint_id') != 1 and inpaint_iter == 1): | |
new_tmp_mask_nodes.append(ne) | |
mask_map[ne[0], ne[1]] = True | |
tmp_mask_nodes = new_tmp_mask_nodes | |
new_tmp_context_nodes = None | |
new_tmp_context_nodes = [] | |
for node in tmp_context_nodes: | |
nes = mesh.neighbors(node) | |
if inpaint_iter > 0: | |
four_nes = {(node[0] - 1, node[1]):[], (node[0] + 1, node[1]):[], (node[0], node[1] - 1):[], (node[0], node[1] + 1):[]} | |
for ne in nes: four_nes[(ne[0], ne[1])].append(ne[2]) | |
nes = [] | |
for kfne, vfnes in four_nes.items(): vfnes.sort(key = lambda xx: abs(xx), reverse=True) | |
for kfne, vfnes in four_nes.items(): | |
for vfne in vfnes: nes.append((kfne[0], kfne[1], vfne)) | |
for ne in nes: | |
mask_flag = (bool(mask_map[ne[0], ne[1]]) is False) | |
if bool(context_map[ne[0], ne[1]]) is False and mask_flag and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True and bool(noncont_map[ne[0], ne[1]]) is False and \ | |
bool(intersect_context_map[ne[0], ne[1]]) is False: | |
if i == passive_context and inpaint_iter == 1: | |
mnes = mesh.neighbors(ne) | |
if any([mask_map[mne[0], mne[1]] == True for mne in mnes]) is True: | |
intersect_context_map[ne[0], ne[1]] = True | |
tmp_intersect_context_nodes.append(ne) | |
continue | |
if False and mesh_nodes[ne].get('near') is not None and mesh_nodes[ne].get('edge_id') != edge_id: | |
noncont_nears = set(mesh_nodes[ne].get('near')) | |
for noncont_near in noncont_nears: | |
if bool(context_map[noncont_near[0], noncont_near[1]]) is False: | |
tmp_noncont_nodes.add(noncont_near) | |
noncont_map[noncont_near[0], noncont_near[1]] = True | |
new_tmp_context_nodes.append(ne) | |
context_map[ne[0], ne[1]] = True | |
context_depth[ne[0], ne[1]] = ne[2] | |
cur_context_cc.extend(new_tmp_context_nodes) | |
tmp_erode.append(new_tmp_context_nodes) | |
tmp_context_nodes = None | |
tmp_context_nodes = new_tmp_context_nodes | |
new_tmp_intouched_nodes = None; new_tmp_intouched_nodes = [] | |
for node in tmp_intouched_nodes: | |
if bool(context_map[node[0], node[1]]) is True or bool(mask_map[node[0], node[1]]) is True: | |
continue | |
nes = mesh.neighbors(node) | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(intouched_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
new_tmp_intouched_nodes.append(ne) | |
intouched_map[ne[0], ne[1]] = True | |
tmp_intouched_nodes = None | |
tmp_intouched_nodes = set(new_tmp_intouched_nodes) | |
new_tmp_redundant_nodes = None; new_tmp_redundant_nodes = [] | |
for node in tmp_redundant_nodes: | |
if bool(context_map[node[0], node[1]]) is True or \ | |
bool(mask_map[node[0], node[1]]) is True: | |
continue | |
nes = mesh.neighbors(node) | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(intouched_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
new_tmp_redundant_nodes.append(ne) | |
intouched_map[ne[0], ne[1]] = True | |
tmp_redundant_nodes = None | |
tmp_redundant_nodes = set(new_tmp_redundant_nodes) | |
new_tmp_noncont_nodes = None; new_tmp_noncont_nodes = [] | |
for node in tmp_noncont_nodes: | |
if bool(context_map[node[0], node[1]]) is True or \ | |
bool(mask_map[node[0], node[1]]) is True: | |
continue | |
nes = mesh.neighbors(node) | |
rmv_flag = False | |
for ne in nes: | |
if bool(context_map[ne[0], ne[1]]) is False and \ | |
bool(mask_map[ne[0], ne[1]]) is False and \ | |
bool(noncont_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
patch_context_map = context_map[max(ne[0] - 1, 0):min(ne[0] + 2, context_map.shape[0]), | |
max(ne[1] - 1, 0):min(ne[1] + 2, context_map.shape[1])] | |
if bool(np.any(patch_context_map)) is True: | |
new_tmp_noncont_nodes.append(ne) | |
noncont_map[ne[0], ne[1]] = True | |
tmp_noncont_nodes = None | |
tmp_noncont_nodes = set(new_tmp_noncont_nodes) | |
if inpaint_iter == 0: | |
depth_dict = get_depth_from_maps(context_map, mask_map, context_depth, mesh.graph['H'], mesh.graph['W'], log_depth=config['log_depth']) | |
mask_size = get_valid_size(depth_dict['mask']) | |
mask_size = dilate_valid_size(mask_size, depth_dict['mask'], dilate=[20, 20]) | |
context_size = get_valid_size(depth_dict['context']) | |
context_size = dilate_valid_size(context_size, depth_dict['context'], dilate=[20, 20]) | |
union_size = size_operation(mask_size, context_size, operation='+') | |
depth_dict = depth_inpainting(None, None, None, None, mesh, config, union_size, depth_feat_model, None, given_depth_dict=depth_dict, spdb=False) | |
near_depth_map, raw_near_depth_map = np.zeros((mesh.graph['H'], mesh.graph['W'])), np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
filtered_comp_far_cc, filtered_accomp_near_cc = set(), set() | |
for node in cur_accomp_near_cc: | |
near_depth_map[node[0], node[1]] = depth_dict['output'][node[0], node[1]] | |
raw_near_depth_map[node[0], node[1]] = node[2] | |
for node in cur_comp_far_cc: | |
four_nes = [xx for xx in [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)] \ | |
if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \ | |
near_depth_map[xx[0], xx[1]] != 0 and \ | |
abs(near_depth_map[xx[0], xx[1]]) < abs(node[2])] | |
if len(four_nes) > 0: | |
filtered_comp_far_cc.add(node) | |
for ne in four_nes: | |
filtered_accomp_near_cc.add((ne[0], ne[1], -abs(raw_near_depth_map[ne[0], ne[1]]))) | |
cur_comp_far_cc, cur_accomp_near_cc = filtered_comp_far_cc, filtered_accomp_near_cc | |
mask_ccs[edge_id] |= set(cur_mask_cc) | |
context_ccs[edge_id] |= set(cur_context_cc) | |
accomp_extend_context_ccs[edge_id] |= set(cur_accomp_near_cc).intersection(cur_mask_cc) | |
extend_edge_ccs[edge_id] |= set(cur_accomp_near_cc).intersection(cur_mask_cc) | |
extend_context_ccs[edge_id] |= set(cur_comp_far_cc) | |
invalid_extend_edge_ccs[edge_id] |= set(cur_invalid_extend_edge_cc) | |
erode_size = [0] | |
for tmp in tmp_erode: | |
erode_size.append(len(tmp)) | |
if len(erode_size) > 1: | |
erode_size[-1] += erode_size[-2] | |
if inpaint_iter == 0: | |
tmp_width = config['depth_edge_dilate'] | |
else: | |
tmp_width = 0 | |
while float(erode_size[tmp_width]) / (erode_size[-1] + 1e-6) > 0.3: | |
tmp_width = tmp_width - 1 | |
try: | |
if tmp_width == 0: | |
erode_context_ccs[edge_id] = set([]) | |
else: | |
erode_context_ccs[edge_id] = set(reduce(lambda x, y : x + y, [] + tmp_erode[:tmp_width])) | |
except: | |
import pdb; pdb.set_trace() | |
erode_context_cc = copy.deepcopy(erode_context_ccs[edge_id]) | |
for erode_context_node in erode_context_cc: | |
if (inpaint_iter != 0 and (mesh_nodes[erode_context_node].get('inpaint_id') is None or | |
mesh_nodes[erode_context_node].get('inpaint_id') == 0)): | |
erode_context_ccs[edge_id].remove(erode_context_node) | |
else: | |
context_ccs[edge_id].remove(erode_context_node) | |
context_map = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
for context_node in context_ccs[edge_id]: | |
context_map[context_node[0], context_node[1]] = 1 | |
extend_context_ccs[edge_id] = extend_context_ccs[edge_id] - mask_ccs[edge_id] - accomp_extend_context_ccs[edge_id] | |
if inpaint_iter == 0: | |
all_ecnt_cc = set() | |
for ecnt_id, ecnt_cc in enumerate(extend_context_ccs): | |
constraint_context_ids = set() | |
constraint_context_cc = set() | |
constraint_erode_context_cc = set() | |
tmp_mask_cc = set() | |
accum_context_cc = None; accum_context_cc = [] | |
for ecnt_node in accomp_extend_context_ccs[ecnt_id]: | |
if edge_maps[ecnt_node[0], ecnt_node[1]] > -1: | |
constraint_context_ids.add(int(round(edge_maps[ecnt_node[0], ecnt_node[1]]))) | |
constraint_erode_context_cc = erode_context_ccs[ecnt_id] | |
for constraint_context_id in constraint_context_ids: | |
constraint_context_cc = constraint_context_cc | context_ccs[constraint_context_id] | erode_context_ccs[constraint_context_id] | |
constraint_erode_context_cc = constraint_erode_context_cc | erode_context_ccs[constraint_context_id] | |
for i in range(background_thickness): | |
if i == 0: | |
tmp_context_nodes = copy.deepcopy(ecnt_cc) | |
tmp_invalid_context_nodes = copy.deepcopy(invalid_extend_edge_ccs[ecnt_id]) | |
tmp_mask_nodes = copy.deepcopy(accomp_extend_context_ccs[ecnt_id]) | |
tmp_context_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool) | |
tmp_mask_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool) | |
tmp_invalid_context_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool) | |
for node in tmp_mask_nodes: | |
tmp_mask_map[node[0], node[1]] = True | |
for node in context_ccs[ecnt_id]: | |
tmp_context_map[node[0], node[1]] = True | |
for node in erode_context_ccs[ecnt_id]: | |
tmp_context_map[node[0], node[1]] = True | |
for node in extend_context_ccs[ecnt_id]: | |
tmp_context_map[node[0], node[1]] = True | |
for node in invalid_extend_edge_ccs[ecnt_id]: | |
tmp_invalid_context_map[node[0], node[1]] = True | |
init_invalid_context_map = tmp_invalid_context_map.copy() | |
init_context_map = tmp | |
if (tmp_mask_map.astype(np.uint8) * tmp_context_map.astype(np.uint8)).max() > 0: | |
import pdb; pdb.set_trace() | |
if vis_edge_id is not None and ecnt_id == vis_edge_id: | |
f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True) | |
ax1.imshow(tmp_context_map * 1); ax2.imshow(init_invalid_context_map * 1 + tmp_context_map * 2) | |
plt.show() | |
import pdb; pdb.set_trace() | |
else: | |
tmp_context_nodes = new_tmp_context_nodes | |
new_tmp_context_nodes = None | |
tmp_mask_nodes = new_tmp_mask_nodes | |
new_tmp_mask_nodes = None | |
tmp_invalid_context_nodes = new_tmp_invalid_context_nodes | |
new_tmp_invalid_context_nodes = None | |
new_tmp_context_nodes = None | |
new_tmp_context_nodes = [] | |
new_tmp_invalid_context_nodes = None | |
new_tmp_invalid_context_nodes = [] | |
new_tmp_mask_nodes = set([]) | |
for node in tmp_context_nodes: | |
for ne in mesh.neighbors(node): | |
if ne in constraint_context_cc and \ | |
bool(tmp_mask_map[ne[0], ne[1]]) is False and \ | |
bool(tmp_context_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
new_tmp_context_nodes.append(ne) | |
tmp_context_map[ne[0], ne[1]] = True | |
accum_context_cc.extend(new_tmp_context_nodes) | |
for node in tmp_invalid_context_nodes: | |
for ne in mesh.neighbors(node): | |
if bool(tmp_mask_map[ne[0], ne[1]]) is False and \ | |
bool(tmp_context_map[ne[0], ne[1]]) is False and \ | |
bool(tmp_invalid_context_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
tmp_invalid_context_map[ne[0], ne[1]] = True | |
new_tmp_invalid_context_nodes.append(ne) | |
for node in tmp_mask_nodes: | |
for ne in mesh.neighbors(node): | |
if bool(tmp_mask_map[ne[0], ne[1]]) is False and \ | |
bool(tmp_context_map[ne[0], ne[1]]) is False and \ | |
bool(tmp_invalid_context_map[ne[0], ne[1]]) is False and \ | |
bool(forbidden_map[ne[0], ne[1]]) is True: | |
new_tmp_mask_nodes.add(ne) | |
tmp_mask_map[ne[0], ne[1]] = True | |
init_invalid_context_map[tmp_context_map] = False | |
_, tmp_label_map = cv2.connectedComponents((init_invalid_context_map | tmp_context_map).astype(np.uint8), connectivity=8) | |
tmp_label_ids = set(np.unique(tmp_label_map[init_invalid_context_map])) | |
if (tmp_mask_map.astype(np.uint8) * tmp_context_map.astype(np.uint8)).max() > 0: | |
import pdb; pdb.set_trace() | |
if vis_edge_id is not None and ecnt_id == vis_edge_id: | |
f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True) | |
ax1.imshow(tmp_label_map); ax2.imshow(init_invalid_context_map * 1 + tmp_context_map * 2) | |
plt.show() | |
import pdb; pdb.set_trace() | |
extend_context_ccs[ecnt_id] |= set(accum_context_cc) | |
extend_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] - mask_ccs[ecnt_id] | |
extend_erode_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] & constraint_erode_context_cc | |
extend_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] - extend_erode_context_ccs[ecnt_id] - erode_context_ccs[ecnt_id] | |
tmp_context_cc = context_ccs[ecnt_id] - extend_erode_context_ccs[ecnt_id] - erode_context_ccs[ecnt_id] | |
if len(tmp_context_cc) > 0: | |
context_ccs[ecnt_id] = tmp_context_cc | |
tmp_mask_cc = tmp_mask_cc - context_ccs[ecnt_id] - erode_context_ccs[ecnt_id] | |
mask_ccs[ecnt_id] = mask_ccs[ecnt_id] | tmp_mask_cc | |
return context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, invalid_extend_edge_ccs, edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs | |
def DL_inpaint_edge(mesh, | |
info_on_pix, | |
config, | |
image, | |
depth, | |
context_ccs, | |
erode_context_ccs, | |
extend_context_ccs, | |
extend_erode_context_ccs, | |
mask_ccs, | |
broken_mask_ccs, | |
edge_ccs, | |
extend_edge_ccs, | |
init_mask_connect, | |
edge_maps, | |
rgb_model=None, | |
depth_edge_model=None, | |
depth_edge_model_init=None, | |
depth_feat_model=None, | |
specific_edge_id=-1, | |
specific_edge_loc=None, | |
inpaint_iter=0): | |
if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0): | |
device = config["gpu_ids"] | |
else: | |
device = "cpu" | |
edge_map = np.zeros_like(depth) | |
new_edge_ccs = [set() for _ in range(len(edge_ccs))] | |
edge_maps_with_id = edge_maps | |
edge_condition = lambda x, m: m.nodes[x].get('far') is not None and len(m.nodes[x].get('far')) > 0 | |
edge_map = get_map_from_ccs(edge_ccs, mesh.graph['H'], mesh.graph['W'], mesh, edge_condition) | |
np_depth, np_image = depth.copy(), image.copy() | |
image_c = image.shape[-1] | |
image = torch.FloatTensor(image.transpose(2, 0, 1)).unsqueeze(0).to(device) | |
if depth.ndim < 3: | |
depth = depth[..., None] | |
depth = torch.FloatTensor(depth.transpose(2, 0, 1)).unsqueeze(0).to(device) | |
mesh.graph['max_edge_id'] = len(edge_ccs) | |
connnect_points_ccs = [set() for _ in range(len(edge_ccs))] | |
gp_time, tmp_mesh_time, bilateral_time = 0, 0, 0 | |
edges_infos = dict() | |
edges_in_mask = [set() for _ in range(len(edge_ccs))] | |
tmp_specific_edge_id = [] | |
for edge_id, (context_cc, mask_cc, erode_context_cc, extend_context_cc, edge_cc) in enumerate(zip(context_ccs, mask_ccs, erode_context_ccs, extend_context_ccs, edge_ccs)): | |
if len(specific_edge_id) > 0: | |
if edge_id not in specific_edge_id: | |
continue | |
if len(context_cc) < 1 or len(mask_cc) < 1: | |
continue | |
edge_dict = get_edge_from_nodes(context_cc | extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, edge_cc, extend_edge_ccs[edge_id], | |
mesh.graph['H'], mesh.graph['W'], mesh) | |
edge_dict['edge'], end_depth_maps, _ = \ | |
filter_irrelevant_edge_new(edge_dict['self_edge'], edge_dict['comp_edge'], | |
edge_map, | |
edge_maps_with_id, | |
edge_id, | |
edge_dict['context'], | |
edge_dict['depth'], mesh, context_cc | erode_context_cc | extend_context_cc | extend_erode_context_ccs[edge_id], spdb=False) | |
if specific_edge_loc is not None and \ | |
(specific_edge_loc is not None and edge_dict['mask'][specific_edge_loc[0], specific_edge_loc[1]] == 0): | |
continue | |
mask_size = get_valid_size(edge_dict['mask']) | |
mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20]) | |
context_size = get_valid_size(edge_dict['context']) | |
context_size = dilate_valid_size(context_size, edge_dict['context'], dilate=[20, 20]) | |
union_size = size_operation(mask_size, context_size, operation='+') | |
patch_edge_dict = dict() | |
patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \ | |
patch_edge_dict['disp'], patch_edge_dict['edge'] = \ | |
crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'], | |
edge_dict['rgb'], edge_dict['disp'], edge_dict['edge']) | |
x_anchor, y_anchor = [union_size['x_min'], union_size['x_max']], [union_size['y_min'], union_size['y_max']] | |
tensor_edge_dict = convert2tensor(patch_edge_dict) | |
input_edge_feat = torch.cat((tensor_edge_dict['rgb'], | |
tensor_edge_dict['disp'], | |
tensor_edge_dict['edge'], | |
1 - tensor_edge_dict['context'], | |
tensor_edge_dict['mask']), dim=1) | |
if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0: | |
with torch.no_grad(): | |
depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'], | |
tensor_edge_dict['context'], | |
tensor_edge_dict['rgb'], | |
tensor_edge_dict['disp'], | |
tensor_edge_dict['edge'], | |
unit_length=128, | |
cuda=device) | |
depth_edge_output = depth_edge_output.cpu() | |
tensor_edge_dict['output'] = (depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge'] | |
else: | |
tensor_edge_dict['output'] = tensor_edge_dict['edge'] | |
depth_edge_output = tensor_edge_dict['edge'] + 0 | |
patch_edge_dict['output'] = tensor_edge_dict['output'].squeeze().data.cpu().numpy() | |
edge_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
edge_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \ | |
patch_edge_dict['output'] | |
if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0: | |
if ((depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask']).max() > 0: | |
try: | |
edge_dict['fpath_map'], edge_dict['npath_map'], break_flag, npaths, fpaths, invalid_edge_id = \ | |
clean_far_edge_new(edge_dict['output'], end_depth_maps, edge_dict['mask'], edge_dict['context'], mesh, info_on_pix, edge_dict['self_edge'], inpaint_iter, config) | |
except: | |
import pdb; pdb.set_trace() | |
pre_npath_map = edge_dict['npath_map'].copy() | |
if config.get('repeat_inpaint_edge') is True: | |
for _ in range(2): | |
tmp_input_edge = ((edge_dict['npath_map'] > -1) + edge_dict['edge']).clip(0, 1) | |
patch_tmp_input_edge = crop_maps_by_size(union_size, tmp_input_edge)[0] | |
tensor_input_edge = torch.FloatTensor(patch_tmp_input_edge)[None, None, ...] | |
depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'], | |
tensor_edge_dict['context'], | |
tensor_edge_dict['rgb'], | |
tensor_edge_dict['disp'], | |
tensor_input_edge, | |
unit_length=128, | |
cuda=device) | |
depth_edge_output = depth_edge_output.cpu() | |
depth_edge_output = (depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge'] | |
depth_edge_output = depth_edge_output.squeeze().data.cpu().numpy() | |
full_depth_edge_output = np.zeros((mesh.graph['H'], mesh.graph['W'])) | |
full_depth_edge_output[union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \ | |
depth_edge_output | |
edge_dict['fpath_map'], edge_dict['npath_map'], break_flag, npaths, fpaths, invalid_edge_id = \ | |
clean_far_edge_new(full_depth_edge_output, end_depth_maps, edge_dict['mask'], edge_dict['context'], mesh, info_on_pix, edge_dict['self_edge'], inpaint_iter, config) | |
for nid in npaths.keys(): | |
npath, fpath = npaths[nid], fpaths[nid] | |
start_mx, start_my, end_mx, end_my = -1, -1, -1, -1 | |
if end_depth_maps[npath[0][0], npath[0][1]] != 0: | |
start_mx, start_my = npath[0][0], npath[0][1] | |
if end_depth_maps[npath[-1][0], npath[-1][1]] != 0: | |
end_mx, end_my = npath[-1][0], npath[-1][1] | |
if start_mx == -1: | |
import pdb; pdb.set_trace() | |
valid_end_pt = () if end_mx == -1 else (end_mx, end_my, info_on_pix[(end_mx, end_my)][0]['depth']) | |
new_edge_info = dict(fpath=fpath, | |
npath=npath, | |
cont_end_pts=valid_end_pt, | |
mask_id=edge_id, | |
comp_edge_id=nid, | |
depth=end_depth_maps[start_mx, start_my]) | |
if edges_infos.get((start_mx, start_my)) is None: | |
edges_infos[(start_mx, start_my)] = [] | |
edges_infos[(start_mx, start_my)].append(new_edge_info) | |
edges_in_mask[edge_id].add((start_mx, start_my)) | |
if len(valid_end_pt) > 0: | |
new_edge_info = dict(fpath=fpath[::-1], | |
npath=npath[::-1], | |
cont_end_pts=(start_mx, start_my, info_on_pix[(start_mx, start_my)][0]['depth']), | |
mask_id=edge_id, | |
comp_edge_id=nid, | |
depth=end_depth_maps[end_mx, end_my]) | |
if edges_infos.get((end_mx, end_my)) is None: | |
edges_infos[(end_mx, end_my)] = [] | |
edges_infos[(end_mx, end_my)].append(new_edge_info) | |
edges_in_mask[edge_id].add((end_mx, end_my)) | |
for edge_id, (context_cc, mask_cc, erode_context_cc, extend_context_cc, edge_cc) in enumerate(zip(context_ccs, mask_ccs, erode_context_ccs, extend_context_ccs, edge_ccs)): | |
if len(specific_edge_id) > 0: | |
if edge_id not in specific_edge_id: | |
continue | |
if len(context_cc) < 1 or len(mask_cc) < 1: | |
continue | |
edge_dict = get_edge_from_nodes(context_cc | extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, edge_cc, extend_edge_ccs[edge_id], | |
mesh.graph['H'], mesh.graph['W'], mesh) | |
if specific_edge_loc is not None and \ | |
(specific_edge_loc is not None and edge_dict['mask'][specific_edge_loc[0], specific_edge_loc[1]] == 0): | |
continue | |
else: | |
tmp_specific_edge_id.append(edge_id) | |
edge_dict['edge'], end_depth_maps, _ = \ | |
filter_irrelevant_edge_new(edge_dict['self_edge'], edge_dict['comp_edge'], | |
edge_map, | |
edge_maps_with_id, | |
edge_id, | |
edge_dict['context'], | |
edge_dict['depth'], mesh, context_cc | erode_context_cc | extend_context_cc | extend_erode_context_ccs[edge_id], spdb=False) | |
discard_map = np.zeros_like(edge_dict['edge']) | |
mask_size = get_valid_size(edge_dict['mask']) | |
mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20]) | |
context_size = get_valid_size(edge_dict['context']) | |
context_size = dilate_valid_size(context_size, edge_dict['context'], dilate=[20, 20]) | |
union_size = size_operation(mask_size, context_size, operation='+') | |
patch_edge_dict = dict() | |
patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \ | |
patch_edge_dict['disp'], patch_edge_dict['edge'] = \ | |
crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'], | |
edge_dict['rgb'], edge_dict['disp'], edge_dict['edge']) | |
x_anchor, y_anchor = [union_size['x_min'], union_size['x_max']], [union_size['y_min'], union_size['y_max']] | |
tensor_edge_dict = convert2tensor(patch_edge_dict) | |
input_edge_feat = torch.cat((tensor_edge_dict['rgb'], | |
tensor_edge_dict['disp'], | |
tensor_edge_dict['edge'], | |
1 - tensor_edge_dict['context'], | |
tensor_edge_dict['mask']), dim=1) | |
edge_dict['output'] = edge_dict['edge'].copy() | |
if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0: | |
edge_dict['fpath_map'], edge_dict['npath_map'] = edge_dict['fpath_map'] * 0 - 1, edge_dict['npath_map'] * 0 - 1 | |
end_pts = edges_in_mask[edge_id] | |
for end_pt in end_pts: | |
cur_edge_infos = edges_infos[(end_pt[0], end_pt[1])] | |
cur_info = [xx for xx in cur_edge_infos if xx['mask_id'] == edge_id][0] | |
other_infos = [xx for xx in cur_edge_infos if xx['mask_id'] != edge_id and len(xx['cont_end_pts']) > 0] | |
if len(cur_info['cont_end_pts']) > 0 or (len(cur_info['cont_end_pts']) == 0 and len(other_infos) == 0): | |
for fnode in cur_info['fpath']: | |
edge_dict['fpath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id'] | |
for fnode in cur_info['npath']: | |
edge_dict['npath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id'] | |
fnmap = edge_dict['fpath_map'] * 1 | |
fnmap[edge_dict['npath_map'] != -1] = edge_dict['npath_map'][edge_dict['npath_map'] != -1] | |
for end_pt in end_pts: | |
cur_edge_infos = edges_infos[(end_pt[0], end_pt[1])] | |
cur_info = [xx for xx in cur_edge_infos if xx['mask_id'] == edge_id][0] | |
cur_depth = cur_info['depth'] | |
other_infos = [xx for xx in cur_edge_infos if xx['mask_id'] != edge_id and len(xx['cont_end_pts']) > 0] | |
comp_edge_id = cur_info['comp_edge_id'] | |
if len(cur_info['cont_end_pts']) == 0 and len(other_infos) > 0: | |
other_infos = sorted(other_infos, key=lambda aa: abs(abs(aa['cont_end_pts'][2]) - abs(cur_depth))) | |
for other_info in other_infos: | |
tmp_fmap, tmp_nmap = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1, np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1 | |
for fnode in other_info['fpath']: | |
if fnmap[fnode[0], fnode[1]] != -1: | |
tmp_fmap = tmp_fmap * 0 - 1 | |
break | |
else: | |
tmp_fmap[fnode[0], fnode[1]] = comp_edge_id | |
if fnmap[fnode[0], fnode[1]] != -1: | |
continue | |
for fnode in other_info['npath']: | |
if fnmap[fnode[0], fnode[1]] != -1: | |
tmp_nmap = tmp_nmap * 0 - 1 | |
break | |
else: | |
tmp_nmap[fnode[0], fnode[1]] = comp_edge_id | |
if fnmap[fnode[0], fnode[1]] != -1: | |
continue | |
break | |
if min(tmp_fmap.max(), tmp_nmap.max()) != -1: | |
edge_dict['fpath_map'] = tmp_fmap | |
edge_dict['fpath_map'][edge_dict['valid_area'] == 0] = -1 | |
edge_dict['npath_map'] = tmp_nmap | |
edge_dict['npath_map'][edge_dict['valid_area'] == 0] = -1 | |
discard_map = ((tmp_nmap != -1).astype(np.uint8) + (tmp_fmap != -1).astype(np.uint8)) * edge_dict['mask'] | |
else: | |
for fnode in cur_info['fpath']: | |
edge_dict['fpath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id'] | |
for fnode in cur_info['npath']: | |
edge_dict['npath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id'] | |
if edge_dict['npath_map'].min() == 0 or edge_dict['fpath_map'].min() == 0: | |
import pdb; pdb.set_trace() | |
edge_dict['output'] = (edge_dict['npath_map'] > -1) * edge_dict['mask'] + edge_dict['context'] * edge_dict['edge'] | |
mesh, _, _, _ = create_placeholder(edge_dict['context'], edge_dict['mask'], | |
edge_dict['depth'], edge_dict['fpath_map'], | |
edge_dict['npath_map'], mesh, inpaint_iter, | |
edge_ccs, | |
extend_edge_ccs[edge_id], | |
edge_maps_with_id, | |
edge_id) | |
dxs, dys = np.where(discard_map != 0) | |
for dx, dy in zip(dxs, dys): | |
mesh.nodes[(dx, dy)]['inpaint_twice'] = False | |
depth_dict = depth_inpainting(context_cc, extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, mesh, config, union_size, depth_feat_model, edge_dict['output']) | |
refine_depth_output = depth_dict['output']*depth_dict['mask'] | |
for near_id in np.unique(edge_dict['npath_map'])[1:]: | |
refine_depth_output = refine_depth_around_edge(refine_depth_output.copy(), | |
(edge_dict['fpath_map'] == near_id).astype(np.uint8) * edge_dict['mask'], | |
(edge_dict['fpath_map'] == near_id).astype(np.uint8), | |
(edge_dict['npath_map'] == near_id).astype(np.uint8) * edge_dict['mask'], | |
depth_dict['mask'].copy(), | |
depth_dict['output'] * depth_dict['context'], | |
config) | |
depth_dict['output'][depth_dict['mask'] > 0] = refine_depth_output[depth_dict['mask'] > 0] | |
rgb_dict = get_rgb_from_nodes(context_cc | extend_context_cc, | |
erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, mesh.graph['H'], mesh.graph['W'], mesh) | |
if np.all(rgb_dict['mask'] == edge_dict['mask']) is False: | |
import pdb; pdb.set_trace() | |
rgb_dict['edge'] = edge_dict['output'] | |
patch_rgb_dict = dict() | |
patch_rgb_dict['mask'], patch_rgb_dict['context'], patch_rgb_dict['rgb'], \ | |
patch_rgb_dict['edge'] = crop_maps_by_size(union_size, rgb_dict['mask'], | |
rgb_dict['context'], rgb_dict['rgb'], | |
rgb_dict['edge']) | |
tensor_rgb_dict = convert2tensor(patch_rgb_dict) | |
resize_rgb_dict = {k: v.clone() for k, v in tensor_rgb_dict.items()} | |
max_hw = np.array([*patch_rgb_dict['mask'].shape[-2:]]).max() | |
init_frac = config['largest_size'] / (np.array([*patch_rgb_dict['mask'].shape[-2:]]).prod() ** 0.5) | |
resize_hw = [patch_rgb_dict['mask'].shape[-2] * init_frac, patch_rgb_dict['mask'].shape[-1] * init_frac] | |
resize_max_hw = max(resize_hw) | |
frac = (np.floor(resize_max_hw / 128.) * 128.) / max_hw | |
if frac < 1: | |
resize_mark = torch.nn.functional.interpolate(torch.cat((resize_rgb_dict['mask'], | |
resize_rgb_dict['context']), | |
dim=1), | |
scale_factor=frac, | |
mode='area') | |
resize_rgb_dict['mask'] = (resize_mark[:, 0:1] > 0).float() | |
resize_rgb_dict['context'] = (resize_mark[:, 1:2] == 1).float() | |
resize_rgb_dict['context'][resize_rgb_dict['mask'] > 0] = 0 | |
resize_rgb_dict['rgb'] = torch.nn.functional.interpolate(resize_rgb_dict['rgb'], | |
scale_factor=frac, | |
mode='area') | |
resize_rgb_dict['rgb'] = resize_rgb_dict['rgb'] * resize_rgb_dict['context'] | |
resize_rgb_dict['edge'] = torch.nn.functional.interpolate(resize_rgb_dict['edge'], | |
scale_factor=frac, | |
mode='area') | |
resize_rgb_dict['edge'] = (resize_rgb_dict['edge'] > 0).float() * 0 | |
resize_rgb_dict['edge'] = resize_rgb_dict['edge'] * (resize_rgb_dict['context'] + resize_rgb_dict['mask']) | |
rgb_input_feat = torch.cat((resize_rgb_dict['rgb'], resize_rgb_dict['edge']), dim=1) | |
rgb_input_feat[:, 3] = 1 - rgb_input_feat[:, 3] | |
resize_mask = open_small_mask(resize_rgb_dict['mask'], resize_rgb_dict['context'], 3, 41) | |
specified_hole = resize_mask | |
with torch.no_grad(): | |
rgb_output = rgb_model.forward_3P(specified_hole, | |
resize_rgb_dict['context'], | |
resize_rgb_dict['rgb'], | |
resize_rgb_dict['edge'], | |
unit_length=128, | |
cuda=device) | |
rgb_output = rgb_output.cpu() | |
if config.get('gray_image') is True: | |
rgb_output = rgb_output.mean(1, keepdim=True).repeat((1,3,1,1)) | |
rgb_output = rgb_output.cpu() | |
resize_rgb_dict['output'] = rgb_output * resize_rgb_dict['mask'] + resize_rgb_dict['rgb'] | |
tensor_rgb_dict['output'] = resize_rgb_dict['output'] | |
if frac < 1: | |
tensor_rgb_dict['output'] = torch.nn.functional.interpolate(tensor_rgb_dict['output'], | |
size=tensor_rgb_dict['mask'].shape[-2:], | |
mode='bicubic') | |
tensor_rgb_dict['output'] = tensor_rgb_dict['output'] * \ | |
tensor_rgb_dict['mask'] + (tensor_rgb_dict['rgb'] * tensor_rgb_dict['context']) | |
patch_rgb_dict['output'] = tensor_rgb_dict['output'].data.cpu().numpy().squeeze().transpose(1,2,0) | |
rgb_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W'], 3)) | |
rgb_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \ | |
patch_rgb_dict['output'] | |
if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) or inpaint_iter > 0: | |
edge_occlusion = True | |
else: | |
edge_occlusion = False | |
for node in erode_context_cc: | |
if rgb_dict['mask'][node[0], node[1]] > 0: | |
for info in info_on_pix[(node[0], node[1])]: | |
if abs(info['depth']) == abs(node[2]): | |
info['update_color'] = (rgb_dict['output'][node[0], node[1]] * 255).astype(np.uint8) | |
if frac < 1.: | |
depth_edge_dilate_2_color_flag = False | |
else: | |
depth_edge_dilate_2_color_flag = True | |
hxs, hys = np.where((rgb_dict['mask'] > 0) & (rgb_dict['erode'] == 0)) | |
for hx, hy in zip(hxs, hys): | |
real_depth = None | |
if abs(depth_dict['output'][hx, hy]) <= abs(np_depth[hx, hy]): | |
depth_dict['output'][hx, hy] = np_depth[hx, hy] + 0.01 | |
node = (hx, hy, -depth_dict['output'][hx, hy]) | |
if info_on_pix.get((node[0], node[1])) is not None: | |
for info in info_on_pix.get((node[0], node[1])): | |
if info.get('inpaint_id') is None or abs(info['inpaint_id'] < mesh.nodes[(hx, hy)]['inpaint_id']): | |
pre_depth = info['depth'] if info.get('real_depth') is None else info['real_depth'] | |
if abs(node[2]) < abs(pre_depth): | |
node = (node[0], node[1], -(abs(pre_depth) + 0.001)) | |
if mesh.has_node(node): | |
real_depth = node[2] | |
while True: | |
if mesh.has_node(node): | |
node = (node[0], node[1], -(abs(node[2]) + 0.001)) | |
else: | |
break | |
if real_depth == node[2]: | |
real_depth = None | |
cur_disp = 1./node[2] | |
if not(mesh.has_node(node)): | |
if not mesh.has_node((node[0], node[1])): | |
print("2D node not found.") | |
import pdb; pdb.set_trace() | |
if inpaint_iter == 1: | |
paint = (rgb_dict['output'][hx, hy] * 255).astype(np.uint8) | |
else: | |
paint = (rgb_dict['output'][hx, hy] * 255).astype(np.uint8) | |
ndict = dict(color=paint, | |
synthesis=True, | |
disp=cur_disp, | |
cc_id=set([edge_id]), | |
overlap_number=1.0, | |
refine_depth=False, | |
edge_occlusion=edge_occlusion, | |
depth_edge_dilate_2_color_flag=depth_edge_dilate_2_color_flag, | |
real_depth=real_depth) | |
mesh, _, _ = refresh_node((node[0], node[1]), mesh.nodes[(node[0], node[1])], node, ndict, mesh, stime=True) | |
if inpaint_iter == 0 and mesh.degree(node) < 4: | |
connnect_points_ccs[edge_id].add(node) | |
if info_on_pix.get((hx, hy)) is None: | |
info_on_pix[(hx, hy)] = [] | |
new_info = {'depth':node[2], | |
'color': paint, | |
'synthesis':True, | |
'disp':cur_disp, | |
'cc_id':set([edge_id]), | |
'inpaint_id':inpaint_iter + 1, | |
'edge_occlusion':edge_occlusion, | |
'overlap_number':1.0, | |
'real_depth': real_depth} | |
info_on_pix[(hx, hy)].append(new_info) | |
specific_edge_id = tmp_specific_edge_id | |
for erode_id, erode_context_cc in enumerate(erode_context_ccs): | |
if len(specific_edge_id) > 0 and erode_id not in specific_edge_id: | |
continue | |
for erode_node in erode_context_cc: | |
for info in info_on_pix[(erode_node[0], erode_node[1])]: | |
if info['depth'] == erode_node[2]: | |
info['color'] = info['update_color'] | |
mesh.nodes[erode_node]['color'] = info['update_color'] | |
np_image[(erode_node[0], erode_node[1])] = info['update_color'] | |
new_edge_ccs = [set() for _ in range(mesh.graph['max_edge_id'] + 1)] | |
for node in mesh.nodes: | |
if len(node) == 2: | |
mesh.remove_node(node) | |
continue | |
if mesh.nodes[node].get('edge_id') is not None and mesh.nodes[node].get('inpaint_id') == inpaint_iter + 1: | |
if mesh.nodes[node].get('inpaint_twice') is False: | |
continue | |
try: | |
new_edge_ccs[mesh.nodes[node].get('edge_id')].add(node) | |
except: | |
import pdb; pdb.set_trace() | |
specific_mask_nodes = None | |
if inpaint_iter == 0: | |
mesh, info_on_pix = refine_color_around_edge(mesh, info_on_pix, new_edge_ccs, config, False) | |
return mesh, info_on_pix, specific_mask_nodes, new_edge_ccs, connnect_points_ccs, np_image | |
def write_ply(image, | |
depth, | |
int_mtx, | |
ply_name, | |
config, | |
rgb_model, | |
depth_edge_model, | |
depth_edge_model_init, | |
depth_feat_model): | |
depth = depth.astype(np.float64) | |
input_mesh, xy2depth, image, depth = create_mesh(depth, image, int_mtx, config) | |
H, W = input_mesh.graph['H'], input_mesh.graph['W'] | |
input_mesh = tear_edges(input_mesh, config['depth_threshold'], xy2depth) | |
input_mesh, info_on_pix = generate_init_node(input_mesh, config, min_node_in_cc=200) | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=False) | |
edge_canvas = np.zeros((H, W)) - 1 | |
input_mesh, info_on_pix, depth = reassign_floating_island(input_mesh, info_on_pix, image, depth) | |
input_mesh = update_status(input_mesh, info_on_pix) | |
specific_edge_id = [] | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
pre_depth = depth.copy() | |
input_mesh, info_on_pix, edge_mesh, depth, aft_mark = remove_dangling(input_mesh, edge_ccs, edge_mesh, info_on_pix, image, depth, config) | |
input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth) | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
edge_canvas = np.zeros((H, W)) - 1 | |
mesh, info_on_pix, depth = fill_missing_node(input_mesh, info_on_pix, image, depth) | |
if config['extrapolate_border'] is True: | |
pre_depth = depth.copy() | |
input_mesh, info_on_pix, depth = refresh_bord_depth(input_mesh, info_on_pix, image, depth) | |
input_mesh = remove_node_feat(input_mesh, 'edge_id') | |
aft_depth = depth.copy() | |
input_mesh, info_on_pix, depth, image = enlarge_border(input_mesh, info_on_pix, depth, image, config) | |
noext_H, noext_W = H, W | |
H, W = image.shape[:2] | |
input_mesh, info_on_pix = fill_dummy_bord(input_mesh, info_on_pix, image, depth, config) | |
edge_ccs, input_mesh, edge_mesh = \ | |
group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth) | |
input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth) | |
edge_ccs, input_mesh, edge_mesh = \ | |
group_edges(input_mesh, config, image, remove_conflict_ordinal=True, spdb=False) | |
input_mesh = remove_redundant_edge(input_mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=config['redundant_number'], spdb=False) | |
input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth) | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth) | |
input_mesh = remove_redundant_edge(input_mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=config['redundant_number'], invalid=True, spdb=False) | |
input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth) | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth) | |
input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth) | |
edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True) | |
edge_condition = lambda x, m: m.nodes[x].get('far') is not None and len(m.nodes[x].get('far')) > 0 | |
edge_map = get_map_from_ccs(edge_ccs, input_mesh.graph['H'], input_mesh.graph['W'], input_mesh, edge_condition) | |
other_edge_with_id = get_map_from_ccs(edge_ccs, input_mesh.graph['H'], input_mesh.graph['W'], real_id=True) | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="up") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="left") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="down") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="right") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="right-up") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="right-down") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="left-up") | |
info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs, | |
depth_edge_model, depth_feat_model, rgb_model, config, direc="left-down") | |
specific_edge_loc = None | |
specific_edge_id = [] | |
vis_edge_id = None | |
context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, \ | |
init_mask_connect, edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs = \ | |
context_and_holes(input_mesh, | |
edge_ccs, | |
config, | |
specific_edge_id, | |
specific_edge_loc, | |
depth_feat_model, | |
inpaint_iter=0, | |
vis_edge_id=vis_edge_id) | |
edge_canvas = np.zeros((H, W)) | |
mask = np.zeros((H, W)) | |
context = np.zeros((H, W)) | |
vis_edge_ccs = filter_edge(input_mesh, edge_ccs, config) | |
edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1 | |
specific_edge_loc = None | |
FG_edge_maps = edge_maps.copy() | |
edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1 | |
# for cc_id, cc in enumerate(edge_ccs): | |
# for node in cc: | |
# edge_canvas[node[0], node[1]] = cc_id | |
# f, ((ax0, ax1, ax2)) = plt.subplots(1, 3, sharex=True, sharey=True); ax0.imshow(1./depth); ax1.imshow(image); ax2.imshow(edge_canvas); plt.show() | |
input_mesh, info_on_pix, specific_edge_nodes, new_edge_ccs, connect_points_ccs, image = DL_inpaint_edge(input_mesh, | |
info_on_pix, | |
config, | |
image, | |
depth, | |
context_ccs, | |
erode_context_ccs, | |
extend_context_ccs, | |
extend_erode_context_ccs, | |
mask_ccs, | |
broken_mask_ccs, | |
edge_ccs, | |
extend_edge_ccs, | |
init_mask_connect, | |
edge_maps, | |
rgb_model, | |
depth_edge_model, | |
depth_edge_model_init, | |
depth_feat_model, | |
specific_edge_id, | |
specific_edge_loc, | |
inpaint_iter=0) | |
specific_edge_id = [] | |
edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) | |
connect_points_ccs = [set() for _ in connect_points_ccs] | |
context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, init_mask_connect, \ | |
edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs = \ | |
context_and_holes(input_mesh, new_edge_ccs, config, specific_edge_id, specific_edge_loc, depth_feat_model, connect_points_ccs, inpaint_iter=1) | |
mask_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) | |
context_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) | |
erode_context_ccs_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) | |
edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) | |
# edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1 | |
# for cc_id, cc in enumerate(edge_ccs): | |
# for node in cc: | |
# edge_canvas[node[0], node[1]] = cc_id | |
specific_edge_id = [] | |
input_mesh, info_on_pix, specific_edge_nodes, new_edge_ccs, _, image = DL_inpaint_edge(input_mesh, | |
info_on_pix, | |
config, | |
image, | |
depth, | |
context_ccs, | |
erode_context_ccs, | |
extend_context_ccs, | |
extend_erode_context_ccs, | |
mask_ccs, | |
broken_mask_ccs, | |
edge_ccs, | |
extend_edge_ccs, | |
init_mask_connect, | |
edge_maps, | |
rgb_model, | |
depth_edge_model, | |
depth_edge_model_init, | |
depth_feat_model, | |
specific_edge_id, | |
specific_edge_loc, | |
inpaint_iter=1) | |
vertex_id = 0 | |
input_mesh.graph['H'], input_mesh.graph['W'] = input_mesh.graph['noext_H'], input_mesh.graph['noext_W'] | |
background_canvas = np.zeros((input_mesh.graph['H'], | |
input_mesh.graph['W'], | |
3)) | |
ply_flag = config.get('save_ply') | |
if ply_flag is True: | |
node_str_list = [] | |
else: | |
node_str_color = [] | |
node_str_point = [] | |
out_fmt = lambda x, x_flag: str(x) if x_flag is True else x | |
point_time = 0 | |
hlight_time = 0 | |
cur_id_time = 0 | |
node_str_time = 0 | |
generate_face_time = 0 | |
point_list = [] | |
k_00, k_02, k_11, k_12 = \ | |
input_mesh.graph['cam_param_pix_inv'][0, 0], input_mesh.graph['cam_param_pix_inv'][0, 2], \ | |
input_mesh.graph['cam_param_pix_inv'][1, 1], input_mesh.graph['cam_param_pix_inv'][1, 2] | |
w_offset = input_mesh.graph['woffset'] | |
h_offset = input_mesh.graph['hoffset'] | |
for pix_xy, pix_list in info_on_pix.items(): | |
for pix_idx, pix_info in enumerate(pix_list): | |
pix_depth = pix_info['depth'] if pix_info.get('real_depth') is None else pix_info['real_depth'] | |
str_pt = [out_fmt(x, ply_flag) for x in reproject_3d_int_detail(pix_xy[0], pix_xy[1], pix_depth, | |
k_00, k_02, k_11, k_12, w_offset, h_offset)] | |
if input_mesh.has_node((pix_xy[0], pix_xy[1], pix_info['depth'])) is False: | |
return False | |
continue | |
if pix_info.get('overlap_number') is not None: | |
str_color = [out_fmt(x, ply_flag) for x in (pix_info['color']/pix_info['overlap_number']).astype(np.uint8).tolist()] | |
else: | |
str_color = [out_fmt(x, ply_flag) for x in pix_info['color'].tolist()] | |
if pix_info.get('edge_occlusion') is True: | |
str_color.append(out_fmt(4, ply_flag)) | |
else: | |
if pix_info.get('inpaint_id') is None: | |
str_color.append(out_fmt(1, ply_flag)) | |
else: | |
str_color.append(out_fmt(pix_info.get('inpaint_id') + 1, ply_flag)) | |
if pix_info.get('modified_border') is True or pix_info.get('ext_pixel') is True: | |
if len(str_color) == 4: | |
str_color[-1] = out_fmt(5, ply_flag) | |
else: | |
str_color.append(out_fmt(5, ply_flag)) | |
pix_info['cur_id'] = vertex_id | |
input_mesh.nodes[(pix_xy[0], pix_xy[1], pix_info['depth'])]['cur_id'] = out_fmt(vertex_id, ply_flag) | |
vertex_id += 1 | |
if ply_flag is True: | |
node_str_list.append(' '.join(str_pt) + ' ' + ' '.join(str_color) + '\n') | |
else: | |
node_str_color.append(str_color) | |
node_str_point.append(str_pt) | |
str_faces = generate_face(input_mesh, info_on_pix, config) | |
if config['save_ply'] is True: | |
print("Writing mesh file %s ..." % ply_name) | |
with open(ply_name, 'w') as ply_fi: | |
ply_fi.write('ply\n' + 'format ascii 1.0\n') | |
ply_fi.write('comment H ' + str(int(input_mesh.graph['H'])) + '\n') | |
ply_fi.write('comment W ' + str(int(input_mesh.graph['W'])) + '\n') | |
ply_fi.write('comment hFov ' + str(float(input_mesh.graph['hFov'])) + '\n') | |
ply_fi.write('comment vFov ' + str(float(input_mesh.graph['vFov'])) + '\n') | |
ply_fi.write('element vertex ' + str(len(node_str_list)) + '\n') | |
ply_fi.write('property float x\n' + \ | |
'property float y\n' + \ | |
'property float z\n' + \ | |
'property uchar red\n' + \ | |
'property uchar green\n' + \ | |
'property uchar blue\n' + \ | |
'property uchar alpha\n') | |
ply_fi.write('element face ' + str(len(str_faces)) + '\n') | |
ply_fi.write('property list uchar int vertex_index\n') | |
ply_fi.write('end_header\n') | |
ply_fi.writelines(node_str_list) | |
ply_fi.writelines(str_faces) | |
ply_fi.close() | |
return input_mesh | |
else: | |
H = int(input_mesh.graph['H']) | |
W = int(input_mesh.graph['W']) | |
hFov = input_mesh.graph['hFov'] | |
vFov = input_mesh.graph['vFov'] | |
node_str_color = np.array(node_str_color).astype(np.float32) | |
node_str_color[..., :3] = node_str_color[..., :3] / 255. | |
node_str_point = np.array(node_str_point) | |
str_faces = np.array(str_faces) | |
return node_str_point, node_str_color, str_faces, H, W, hFov, vFov | |
def read_ply(mesh_fi): | |
ply_fi = open(mesh_fi, 'r') | |
Height = None | |
Width = None | |
hFov = None | |
vFov = None | |
while True: | |
line = ply_fi.readline().split('\n')[0] | |
if line.startswith('element vertex'): | |
num_vertex = int(line.split(' ')[-1]) | |
elif line.startswith('element face'): | |
num_face = int(line.split(' ')[-1]) | |
elif line.startswith('comment'): | |
if line.split(' ')[1] == 'H': | |
Height = int(line.split(' ')[-1].split('\n')[0]) | |
if line.split(' ')[1] == 'W': | |
Width = int(line.split(' ')[-1].split('\n')[0]) | |
if line.split(' ')[1] == 'hFov': | |
hFov = float(line.split(' ')[-1].split('\n')[0]) | |
if line.split(' ')[1] == 'vFov': | |
vFov = float(line.split(' ')[-1].split('\n')[0]) | |
elif line.startswith('end_header'): | |
break | |
contents = ply_fi.readlines() | |
vertex_infos = contents[:num_vertex] | |
face_infos = contents[num_vertex:] | |
verts = [] | |
colors = [] | |
faces = [] | |
for v_info in vertex_infos: | |
str_info = [float(v) for v in v_info.split('\n')[0].split(' ')] | |
if len(str_info) == 6: | |
vx, vy, vz, r, g, b = str_info | |
else: | |
vx, vy, vz, r, g, b, hi = str_info | |
verts.append([vx, vy, vz]) | |
colors.append([r, g, b, hi]) | |
verts = np.array(verts) | |
try: | |
colors = np.array(colors) | |
colors[..., :3] = colors[..., :3]/255. | |
except: | |
import pdb | |
pdb.set_trace() | |
for f_info in face_infos: | |
_, v1, v2, v3 = [int(f) for f in f_info.split('\n')[0].split(' ')] | |
faces.append([v1, v2, v3]) | |
faces = np.array(faces) | |
return verts, colors, faces, Height, Width, hFov, vFov | |
class Canvas_view(): | |
def __init__(self, | |
fov, | |
verts, | |
faces, | |
colors, | |
canvas_size, | |
factor=1, | |
bgcolor='gray', | |
proj='perspective', | |
): | |
self.canvas = scene.SceneCanvas(bgcolor=bgcolor, size=(canvas_size*factor, canvas_size*factor)) | |
self.view = self.canvas.central_widget.add_view() | |
self.view.camera = 'perspective' | |
self.view.camera.fov = fov | |
self.mesh = visuals.Mesh(shading=None) | |
self.mesh.attach(Alpha(1.0)) | |
self.view.add(self.mesh) | |
self.tr = self.view.camera.transform | |
self.mesh.set_data(vertices=verts, faces=faces, vertex_colors=colors[:, :3]) | |
self.translate([0,0,0]) | |
self.rotate(axis=[1,0,0], angle=180) | |
self.view_changed() | |
def translate(self, trans=[0,0,0]): | |
self.tr.translate(trans) | |
def rotate(self, axis=[1,0,0], angle=0): | |
self.tr.rotate(axis=axis, angle=angle) | |
def view_changed(self): | |
self.view.camera.view_changed() | |
def render(self): | |
return self.canvas.render() | |
def reinit_mesh(self, verts, faces, colors): | |
self.mesh.set_data(vertices=verts, faces=faces, vertex_colors=colors[:, :3]) | |
def reinit_camera(self, fov): | |
self.view.camera.fov = fov | |
self.view.camera.view_changed() | |
def output_3d_photo(verts, colors, faces, Height, Width, hFov, vFov, tgt_poses, video_traj_types, ref_pose, | |
output_dir, ref_image, int_mtx, config, image, videos_poses, video_basename, original_H=None, original_W=None, | |
border=None, depth=None, normal_canvas=None, all_canvas=None, mean_loc_depth=None): | |
cam_mesh = netx.Graph() | |
cam_mesh.graph['H'] = Height | |
cam_mesh.graph['W'] = Width | |
cam_mesh.graph['original_H'] = original_H | |
cam_mesh.graph['original_W'] = original_W | |
int_mtx_real_x = int_mtx[0] * Width | |
int_mtx_real_y = int_mtx[1] * Height | |
cam_mesh.graph['hFov'] = 2 * np.arctan((1. / 2.) * ((cam_mesh.graph['original_W']) / int_mtx_real_x[0])) | |
cam_mesh.graph['vFov'] = 2 * np.arctan((1. / 2.) * ((cam_mesh.graph['original_H']) / int_mtx_real_y[1])) | |
colors = colors[..., :3] | |
fov_in_rad = max(cam_mesh.graph['vFov'], cam_mesh.graph['hFov']) | |
fov = (fov_in_rad * 180 / np.pi) | |
print("fov: " + str(fov)) | |
init_factor = 1 | |
if config.get('anti_flickering') is True: | |
init_factor = 3 | |
if (cam_mesh.graph['original_H'] is not None) and (cam_mesh.graph['original_W'] is not None): | |
canvas_w = cam_mesh.graph['original_W'] | |
canvas_h = cam_mesh.graph['original_H'] | |
else: | |
canvas_w = cam_mesh.graph['W'] | |
canvas_h = cam_mesh.graph['H'] | |
canvas_size = max(canvas_h, canvas_w) | |
if normal_canvas is None: | |
normal_canvas = Canvas_view(fov, | |
verts, | |
faces, | |
colors, | |
canvas_size=canvas_size, | |
factor=init_factor, | |
bgcolor='gray', | |
proj='perspective') | |
else: | |
normal_canvas.reinit_mesh(verts, faces, colors) | |
normal_canvas.reinit_camera(fov) | |
img = normal_canvas.render() | |
backup_img, backup_all_img, all_img_wo_bound = img.copy(), img.copy() * 0, img.copy() * 0 | |
img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_AREA) | |
if border is None: | |
border = [0, img.shape[0], 0, img.shape[1]] | |
H, W = cam_mesh.graph['H'], cam_mesh.graph['W'] | |
if (cam_mesh.graph['original_H'] is not None) and (cam_mesh.graph['original_W'] is not None): | |
aspect_ratio = cam_mesh.graph['original_H'] / cam_mesh.graph['original_W'] | |
else: | |
aspect_ratio = cam_mesh.graph['H'] / cam_mesh.graph['W'] | |
if aspect_ratio > 1: | |
img_h_len = cam_mesh.graph['H'] if cam_mesh.graph.get('original_H') is None else cam_mesh.graph['original_H'] | |
img_w_len = img_h_len / aspect_ratio | |
anchor = [0, | |
img.shape[0], | |
int(max(0, int((img.shape[1])//2 - img_w_len//2))), | |
int(min(int((img.shape[1])//2 + img_w_len//2), (img.shape[1])-1))] | |
elif aspect_ratio <= 1: | |
img_w_len = cam_mesh.graph['W'] if cam_mesh.graph.get('original_W') is None else cam_mesh.graph['original_W'] | |
img_h_len = img_w_len * aspect_ratio | |
anchor = [int(max(0, int((img.shape[0])//2 - img_h_len//2))), | |
int(min(int((img.shape[0])//2 + img_h_len//2), (img.shape[0])-1)), | |
0, | |
img.shape[1]] | |
anchor = np.array(anchor) | |
plane_width = np.tan(fov_in_rad/2.) * np.abs(mean_loc_depth) | |
for video_pose, video_traj_type in zip(videos_poses, video_traj_types): | |
stereos = [] | |
tops = []; buttoms = []; lefts = []; rights = [] | |
for tp_id, tp in enumerate(video_pose): | |
rel_pose = np.linalg.inv(np.dot(tp, np.linalg.inv(ref_pose))) | |
axis, angle = transforms3d.axangles.mat2axangle(rel_pose[0:3, 0:3]) | |
normal_canvas.rotate(axis=axis, angle=(angle*180)/np.pi) | |
normal_canvas.translate(rel_pose[:3,3]) | |
new_mean_loc_depth = mean_loc_depth - float(rel_pose[2, 3]) | |
if 'dolly' in video_traj_type: | |
new_fov = float((np.arctan2(plane_width, np.array([np.abs(new_mean_loc_depth)])) * 180. / np.pi) * 2) | |
normal_canvas.reinit_camera(new_fov) | |
else: | |
normal_canvas.reinit_camera(fov) | |
normal_canvas.view_changed() | |
img = normal_canvas.render() | |
img = cv2.GaussianBlur(img,(int(init_factor//2 * 2 + 1), int(init_factor//2 * 2 + 1)), 0) | |
img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_AREA) | |
img = img[anchor[0]:anchor[1], anchor[2]:anchor[3]] | |
img = img[int(border[0]):int(border[1]), int(border[2]):int(border[3])] | |
if any(np.array(config['crop_border']) > 0.0): | |
H_c, W_c, _ = img.shape | |
o_t = int(H_c * config['crop_border'][0]) | |
o_l = int(W_c * config['crop_border'][1]) | |
o_b = int(H_c * config['crop_border'][2]) | |
o_r = int(W_c * config['crop_border'][3]) | |
img = img[o_t:H_c-o_b, o_l:W_c-o_r] | |
img = cv2.resize(img, (W_c, H_c), interpolation=cv2.INTER_CUBIC) | |
""" | |
img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_CUBIC) | |
img = img[anchor[0]:anchor[1], anchor[2]:anchor[3]] | |
img = img[int(border[0]):int(border[1]), int(border[2]):int(border[3])] | |
if config['crop_border'] is True: | |
top, buttom, left, right = find_largest_rect(img, bg_color=(128, 128, 128)) | |
tops.append(top); buttoms.append(buttom); lefts.append(left); rights.append(right) | |
""" | |
stereos.append(img[..., :3]) | |
normal_canvas.translate(-rel_pose[:3,3]) | |
normal_canvas.rotate(axis=axis, angle=-(angle*180)/np.pi) | |
normal_canvas.view_changed() | |
""" | |
if config['crop_border'] is True: | |
atop, abuttom = min(max(tops), img.shape[0]//2 - 10), max(min(buttoms), img.shape[0]//2 + 10) | |
aleft, aright = min(max(lefts), img.shape[1]//2 - 10), max(min(rights), img.shape[1]//2 + 10) | |
atop -= atop % 2; abuttom -= abuttom % 2; aleft -= aleft % 2; aright -= aright % 2 | |
else: | |
atop = 0; abuttom = img.shape[0] - img.shape[0] % 2; aleft = 0; aright = img.shape[1] - img.shape[1] % 2 | |
""" | |
atop = 0; abuttom = img.shape[0] - img.shape[0] % 2; aleft = 0; aright = img.shape[1] - img.shape[1] % 2 | |
crop_stereos = [] | |
for stereo in stereos: | |
crop_stereos.append((stereo[atop:abuttom, aleft:aright, :3] * 1).astype(np.uint8)) | |
stereos = crop_stereos | |
clip = ImageSequenceClip(stereos, fps=config['fps']) | |
if isinstance(video_basename, list): | |
video_basename = video_basename[0] | |
clip.write_videofile(os.path.join(output_dir, video_basename + '_' + video_traj_type + '.mp4'), fps=config['fps']) | |
return normal_canvas, all_canvas | |