2298 lines
134 KiB
Python
2298 lines
134 KiB
Python
![]() |
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 Inpainting.utils import path_planning, open_small_mask, clean_far_edge, refine_depth_around_edge
|
||
|
from Inpainting.utils import refine_color_around_edge, filter_irrelevant_edge_new, require_depth_edge, clean_far_edge_new
|
||
|
from Inpainting.utils import create_placeholder, refresh_node, find_largest_rect
|
||
|
from Inpainting.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 Inpainting.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
|