3754 lines
148 KiB
Python
3754 lines
148 KiB
Python
import os
|
|
import numpy as np
|
|
|
|
try:
|
|
import networkx 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.0]])
|
|
(
|
|
LDI.graph["cam_param_pix"],
|
|
LDI.graph["cam_param_pix_inv"],
|
|
) = int_mtx_pix, np.linalg.inv(int_mtx_pix)
|
|
disp = 1.0 / (-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.0 / (2 * k[0, 0]))
|
|
mesh.graph["vFov"] = 2 * np.arctan(1.0 / (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.0)
|
|
if mesh.graph.get("tan_vFov") is None:
|
|
mesh.graph["tan_vFov"] = np.tan(mesh.graph["vFov"] / 2.0)
|
|
|
|
ray = np.array(
|
|
[
|
|
(-1.0 + 2.0 * ((sy + 0.5 - w_offset) / (mesh.graph["W"] - 1)))
|
|
* mesh.graph["tan_hFov"],
|
|
(1.0 - 2.0 * (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.0 / 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.0 / 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.0 / 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.0 / node[2],
|
|
1.0 / 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.0 / new_depth
|
|
new_node = (hx, hy, new_depth)
|
|
mesh = refresh_node(
|
|
single_edge_node, mesh.nodes[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.0
|
|
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.0
|
|
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.0
|
|
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.0
|
|
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
|
|
passive_context = 1
|
|
|
|
if background_thickness != None:
|
|
passive_background = background_thickness
|
|
if context_thickness != None:
|
|
passive_context = 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.0) * 128.0) / 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.0:
|
|
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.0 / 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.0
|
|
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.0
|
|
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.0 / 2.0) * ((cam_mesh.graph["original_W"]) / int_mtx_real_x[0])
|
|
)
|
|
cam_mesh.graph["vFov"] = 2 * np.arctan(
|
|
(1.0 / 2.0) * ((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.0) * 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.0
|
|
/ 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
|