|
import os |
|
import os.path as osp |
|
|
|
import cupy as cp |
|
import cv2 |
|
import numpy as np |
|
import torch |
|
import trimesh |
|
from cupyx.scipy.sparse import ( |
|
coo_matrix, |
|
csr_matrix, |
|
diags, |
|
hstack, |
|
spdiags, |
|
vstack, |
|
) |
|
from cupyx.scipy.sparse.linalg import cg |
|
from PIL import Image |
|
from tqdm.auto import tqdm |
|
|
|
from lib.dataset.mesh_util import clean_floats |
|
|
|
|
|
def find_max_list(lst): |
|
list_len = [len(i) for i in lst] |
|
max_id = np.argmax(np.array(list_len)) |
|
return lst[max_id] |
|
|
|
|
|
def interpolate_pts(pts, diff_ids): |
|
|
|
pts_extend = np.around((pts[diff_ids] + pts[diff_ids - 1]) * 0.5).astype(np.int32) |
|
pts = np.insert(pts, diff_ids, pts_extend, axis=0) |
|
|
|
return pts |
|
|
|
|
|
def align_pts(pts1, pts2): |
|
|
|
diff_num = abs(len(pts1) - len(pts2)) |
|
diff_ids = np.sort(np.random.choice(min(len(pts2), len(pts1)), diff_num, replace=True)) |
|
|
|
if len(pts1) > len(pts2): |
|
pts2 = interpolate_pts(pts2, diff_ids) |
|
elif len(pts2) > len(pts1): |
|
pts1 = interpolate_pts(pts1, diff_ids) |
|
else: |
|
pass |
|
|
|
return pts1, pts2 |
|
|
|
|
|
def repeat_pts(pts1, pts2): |
|
|
|
coverage_mask = ((pts1[:, None, :] == pts2[None, :, :]).sum(axis=2) == 2.).any(axis=1) |
|
|
|
return coverage_mask |
|
|
|
|
|
def find_contour(mask, method='all'): |
|
|
|
if method == 'all': |
|
|
|
contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) |
|
else: |
|
contours, _ = cv2.findContours( |
|
mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE |
|
) |
|
|
|
contour_cloth = np.array(find_max_list(contours))[:, 0, :] |
|
|
|
return contour_cloth |
|
|
|
|
|
def mean_value_cordinates(inner_pts, contour_pts): |
|
|
|
body_edges_a = np.sqrt(((inner_pts[:, None] - contour_pts[None, :])**2).sum(axis=2)) |
|
body_edges_c = np.roll(body_edges_a, shift=-1, axis=1) |
|
body_edges_b = np.sqrt(((contour_pts - np.roll(contour_pts, shift=-1, axis=0))**2).sum(axis=1)) |
|
|
|
body_edges = np.concatenate([ |
|
body_edges_a[..., None], body_edges_c[..., None], |
|
np.repeat(body_edges_b[None, :, None], axis=0, repeats=len(inner_pts)) |
|
], |
|
axis=-1) |
|
|
|
body_cos = (body_edges[:, :, 0]**2 + body_edges[:, :, 1]**2 - |
|
body_edges[:, :, 2]**2) / (2 * body_edges[:, :, 0] * body_edges[:, :, 1]) |
|
body_tan_half = np.sqrt( |
|
(1. - np.clip(body_cos, a_max=1., a_min=-1.)) / np.clip(1. + body_cos, 1e-6, 2.) |
|
) |
|
|
|
w = (body_tan_half + np.roll(body_tan_half, shift=1, axis=1)) / body_edges_a |
|
w /= w.sum(axis=1, keepdims=True) |
|
|
|
return w |
|
|
|
|
|
def get_dst_mat(contour_body, contour_cloth): |
|
|
|
dst_mat = ((contour_body[:, None, :] - contour_cloth[None, :, :])**2).sum(axis=2) |
|
|
|
return dst_mat |
|
|
|
|
|
def dispCorres(img_size, contour1, contour2, phi, dir_path): |
|
|
|
contour1 = contour1[None, :, None, :].astype(np.int32) |
|
contour2 = contour2[None, :, None, :].astype(np.int32) |
|
|
|
disp = np.zeros((img_size, img_size, 3), dtype=np.uint8) |
|
cv2.drawContours(disp, contour1, -1, (0, 255, 0), 1) |
|
cv2.drawContours(disp, contour2, -1, (255, 0, 0), 1) |
|
|
|
for i in range(contour1.shape[1]): |
|
|
|
|
|
corresPoint = contour2[0, phi[i], 0] |
|
|
|
cv2.line( |
|
disp, (contour1[0, i, 0, 0], contour1[0, i, 0, 1]), (corresPoint[0], corresPoint[1]), |
|
(255, 255, 255), 1 |
|
) |
|
|
|
cv2.imwrite(osp.join(dir_path, "corres.png"), disp) |
|
|
|
|
|
def remove_stretched_faces(verts, faces): |
|
|
|
mesh = trimesh.Trimesh(verts, faces) |
|
camera_ray = np.array([0.0, 0.0, 1.0]) |
|
faces_cam_angles = np.dot(mesh.face_normals, camera_ray) |
|
|
|
|
|
faces_mask = np.abs(faces_cam_angles) > 2e-1 |
|
|
|
mesh.update_faces(faces_mask) |
|
mesh.remove_unreferenced_vertices() |
|
|
|
return mesh.vertices, mesh.faces |
|
|
|
|
|
def tensor2arr(t, mask=False): |
|
if not mask: |
|
return t.squeeze(0).permute(1, 2, 0).detach().cpu().numpy() |
|
else: |
|
mask = t.squeeze(0).abs().sum(dim=0, keepdim=True) |
|
return (mask != mask[:, 0, 0]).float().squeeze(0).detach().cpu().numpy() |
|
|
|
|
|
def arr2png(t): |
|
return ((t + 1.0) * 0.5 * 255.0).astype(np.uint8) |
|
|
|
|
|
def depth2arr(t): |
|
|
|
return t.float().detach().cpu().numpy() |
|
|
|
|
|
def depth2png(t): |
|
|
|
t_copy = t.copy() |
|
t_bg = t_copy[0, 0] |
|
valid_region = np.logical_and(t > -1.0, t != t_bg) |
|
t_copy[valid_region] -= t_copy[valid_region].min() |
|
t_copy[valid_region] /= t_copy[valid_region].max() |
|
t_copy[valid_region] = (1. - t_copy[valid_region]) * 255.0 |
|
t_copy[~valid_region] = 0.0 |
|
|
|
return t_copy[..., None].astype(np.uint8) |
|
|
|
|
|
def verts_transform(t, depth_scale): |
|
|
|
t_copy = t.clone() |
|
t_copy *= depth_scale * 0.5 |
|
t_copy += depth_scale * 0.5 |
|
t_copy = t_copy[:, [1, 0, 2]] * torch.Tensor([2.0, 2.0, -2.0]) + torch.Tensor([ |
|
0.0, 0.0, depth_scale |
|
]) |
|
|
|
return t_copy |
|
|
|
|
|
def verts_inverse_transform(t, depth_scale): |
|
|
|
t_copy = t.clone() |
|
t_copy -= torch.tensor([0.0, 0.0, depth_scale]) |
|
t_copy /= torch.tensor([2.0, 2.0, -2.0]) |
|
t_copy -= depth_scale * 0.5 |
|
t_copy /= depth_scale * 0.5 |
|
t_copy = t_copy[:, [1, 0, 2]] |
|
|
|
return t_copy |
|
|
|
|
|
def depth_inverse_transform(t, depth_scale): |
|
|
|
t_copy = t.clone() |
|
t_copy -= torch.tensor(depth_scale) |
|
t_copy /= torch.tensor(-2.0) |
|
t_copy -= depth_scale * 0.5 |
|
t_copy /= depth_scale * 0.5 |
|
|
|
return t_copy |
|
|
|
|
|
|
|
|
|
|
|
def move_left(mask): |
|
return cp.pad(mask, ((0, 0), (0, 1)), "constant", constant_values=0)[:, 1:] |
|
|
|
|
|
def move_right(mask): |
|
return cp.pad(mask, ((0, 0), (1, 0)), "constant", constant_values=0)[:, :-1] |
|
|
|
|
|
def move_top(mask): |
|
return cp.pad(mask, ((0, 1), (0, 0)), "constant", constant_values=0)[1:, :] |
|
|
|
|
|
def move_bottom(mask): |
|
return cp.pad(mask, ((1, 0), (0, 0)), "constant", constant_values=0)[:-1, :] |
|
|
|
|
|
def move_top_left(mask): |
|
return cp.pad(mask, ((0, 1), (0, 1)), "constant", constant_values=0)[1:, 1:] |
|
|
|
|
|
def move_top_right(mask): |
|
return cp.pad(mask, ((0, 1), (1, 0)), "constant", constant_values=0)[1:, :-1] |
|
|
|
|
|
def move_bottom_left(mask): |
|
return cp.pad(mask, ((1, 0), (0, 1)), "constant", constant_values=0)[:-1, 1:] |
|
|
|
|
|
def move_bottom_right(mask): |
|
return cp.pad(mask, ((1, 0), (1, 0)), "constant", constant_values=0)[:-1, :-1] |
|
|
|
|
|
def generate_dx_dy_new(mask, nz_horizontal, nz_vertical, step_size=1): |
|
|
|
|
|
|
|
|
|
|
|
|
|
num_pixel = cp.sum(mask) |
|
|
|
pixel_idx = cp.zeros_like(mask, dtype=int) |
|
pixel_idx[mask] = cp.arange(num_pixel) |
|
|
|
has_left_mask = cp.logical_and(move_right(mask), mask) |
|
has_right_mask = cp.logical_and(move_left(mask), mask) |
|
has_bottom_mask = cp.logical_and(move_top(mask), mask) |
|
has_top_mask = cp.logical_and(move_bottom(mask), mask) |
|
|
|
nz_left = nz_horizontal[has_left_mask[mask]] |
|
nz_right = nz_horizontal[has_right_mask[mask]] |
|
nz_top = nz_vertical[has_top_mask[mask]] |
|
nz_bottom = nz_vertical[has_bottom_mask[mask]] |
|
|
|
data = cp.stack([-nz_left / step_size, nz_left / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[move_left(has_left_mask)], pixel_idx[has_left_mask]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_left_mask[mask].astype(int) * 2)]) |
|
D_horizontal_neg = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_right / step_size, nz_right / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[has_right_mask], pixel_idx[move_right(has_right_mask)]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_right_mask[mask].astype(int) * 2)]) |
|
D_horizontal_pos = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_top / step_size, nz_top / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[has_top_mask], pixel_idx[move_top(has_top_mask)]), -1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_top_mask[mask].astype(int) * 2)]) |
|
D_vertical_pos = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_bottom / step_size, nz_bottom / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[move_bottom(has_bottom_mask)], pixel_idx[has_bottom_mask]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_bottom_mask[mask].astype(int) * 2)]) |
|
D_vertical_neg = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
return D_horizontal_pos, D_horizontal_neg, D_vertical_pos, D_vertical_neg |
|
|
|
|
|
def generate_dx_dy(mask, nz_horizontal, nz_vertical, step_size=1): |
|
|
|
|
|
|
|
|
|
|
|
|
|
num_pixel = cp.sum(mask) |
|
|
|
pixel_idx = cp.zeros_like(mask, dtype=int) |
|
pixel_idx[mask] = cp.arange(num_pixel) |
|
|
|
has_left_mask = cp.logical_and(move_right(mask), mask) |
|
has_right_mask = cp.logical_and(move_left(mask), mask) |
|
has_bottom_mask = cp.logical_and(move_top(mask), mask) |
|
has_top_mask = cp.logical_and(move_bottom(mask), mask) |
|
|
|
nz_left = nz_horizontal[has_left_mask[mask]] |
|
nz_right = nz_horizontal[has_right_mask[mask]] |
|
nz_top = nz_vertical[has_top_mask[mask]] |
|
nz_bottom = nz_vertical[has_bottom_mask[mask]] |
|
|
|
data = cp.stack([-nz_left / step_size, nz_left / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[move_left(has_left_mask)], pixel_idx[has_left_mask]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_left_mask[mask].astype(int) * 2)]) |
|
D_horizontal_neg = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_right / step_size, nz_right / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[has_right_mask], pixel_idx[move_right(has_right_mask)]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_right_mask[mask].astype(int) * 2)]) |
|
D_horizontal_pos = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_top / step_size, nz_top / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[has_top_mask], pixel_idx[move_top(has_top_mask)]), -1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_top_mask[mask].astype(int) * 2)]) |
|
D_vertical_pos = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
data = cp.stack([-nz_bottom / step_size, nz_bottom / step_size], -1).flatten() |
|
indices = cp.stack((pixel_idx[move_bottom(has_bottom_mask)], pixel_idx[has_bottom_mask]), |
|
-1).flatten() |
|
indptr = cp.concatenate([cp.array([0]), cp.cumsum(has_bottom_mask[mask].astype(int) * 2)]) |
|
D_vertical_neg = csr_matrix((data, indices, indptr), shape=(num_pixel, num_pixel)) |
|
|
|
return D_horizontal_pos, D_horizontal_neg, D_vertical_pos, D_vertical_neg |
|
|
|
|
|
def construct_facets_from(mask): |
|
idx = cp.zeros_like(mask, dtype=int) |
|
idx[mask] = cp.arange(cp.sum(mask)) |
|
|
|
facet_move_top_mask = move_top(mask) |
|
facet_move_left_mask = move_left(mask) |
|
facet_move_top_left_mask = move_top_left(mask) |
|
facet_top_left_mask = ( |
|
facet_move_top_mask * facet_move_left_mask * facet_move_top_left_mask * mask |
|
) |
|
facet_top_right_mask = move_right(facet_top_left_mask) |
|
facet_bottom_left_mask = move_bottom(facet_top_left_mask) |
|
facet_bottom_right_mask = move_bottom_right(facet_top_left_mask) |
|
|
|
return cp.hstack(( |
|
4 * cp.ones((cp.sum(facet_top_left_mask).item(), 1)), |
|
idx[facet_top_left_mask][:, None], |
|
idx[facet_bottom_left_mask][:, None], |
|
idx[facet_bottom_right_mask][:, None], |
|
idx[facet_top_right_mask][:, None], |
|
)).astype(int) |
|
|
|
|
|
def map_depth_map_to_point_clouds(depth_map, mask, K=None, step_size=1): |
|
|
|
|
|
|
|
|
|
|
|
H, W = mask.shape |
|
yy, xx = cp.meshgrid(cp.arange(W), cp.arange(H)) |
|
xx = cp.flip(xx, axis=0) |
|
|
|
if K is None: |
|
vertices = cp.zeros((H, W, 3)) |
|
vertices[..., 0] = xx * step_size |
|
vertices[..., 1] = yy * step_size |
|
vertices[..., 2] = depth_map |
|
vertices = vertices[mask] |
|
else: |
|
u = cp.zeros((H, W, 3)) |
|
u[..., 0] = xx |
|
u[..., 1] = yy |
|
u[..., 2] = 1 |
|
u = u[mask].T |
|
vertices = (cp.linalg.inv(K) @ u).T * depth_map[mask, cp.newaxis] |
|
|
|
return vertices |
|
|
|
|
|
def sigmoid(x, k=1): |
|
return 1 / (1 + cp.exp(-k * x)) |
|
|
|
|
|
def boundary_excluded_mask(mask): |
|
top_mask = cp.pad(mask, ((1, 0), (0, 0)), "constant", constant_values=0)[:-1, :] |
|
bottom_mask = cp.pad(mask, ((0, 1), (0, 0)), "constant", constant_values=0)[1:, :] |
|
left_mask = cp.pad(mask, ((0, 0), (1, 0)), "constant", constant_values=0)[:, :-1] |
|
right_mask = cp.pad(mask, ((0, 0), (0, 1)), "constant", constant_values=0)[:, 1:] |
|
be_mask = top_mask * bottom_mask * left_mask * right_mask * mask |
|
|
|
|
|
top_mask = cp.pad(be_mask, ((1, 0), (0, 0)), "constant", constant_values=0)[:-1, :] |
|
bottom_mask = cp.pad(be_mask, ((0, 1), (0, 0)), "constant", constant_values=0)[1:, :] |
|
left_mask = cp.pad(be_mask, ((0, 0), (1, 0)), "constant", constant_values=0)[:, :-1] |
|
right_mask = cp.pad(be_mask, ((0, 0), (0, 1)), "constant", constant_values=0)[:, 1:] |
|
bes_mask = (top_mask + bottom_mask + left_mask + right_mask).astype(bool) |
|
be_mask = cp.logical_and(be_mask, bes_mask) |
|
return be_mask |
|
|
|
|
|
def create_boundary_matrix(mask): |
|
num_pixel = cp.sum(mask) |
|
pixel_idx = cp.zeros_like(mask, dtype=int) |
|
pixel_idx[mask] = cp.arange(num_pixel) |
|
|
|
be_mask = boundary_excluded_mask(mask) |
|
boundary_mask = cp.logical_xor(be_mask, mask) |
|
diag_data_term = boundary_mask[mask].astype(int) |
|
B = diags(diag_data_term) |
|
|
|
num_boundary_pixel = cp.sum(boundary_mask).item() |
|
data_term = cp.concatenate((cp.ones(num_boundary_pixel), -cp.ones(num_boundary_pixel))) |
|
row_idx = cp.tile(cp.arange(num_boundary_pixel), 2) |
|
col_idx = cp.concatenate((pixel_idx[boundary_mask], pixel_idx[boundary_mask] + num_pixel)) |
|
B_full = coo_matrix((data_term, (row_idx, col_idx)), shape=(num_boundary_pixel, 2 * num_pixel)) |
|
return B, B_full |
|
|
|
|
|
def double_side_bilateral_normal_integration( |
|
normal_front, |
|
normal_back, |
|
normal_mask, |
|
depth_front=None, |
|
depth_back=None, |
|
depth_mask=None, |
|
k=2, |
|
lambda_normal_back=1, |
|
lambda_depth_front=1e-4, |
|
lambda_depth_back=1e-2, |
|
lambda_boundary_consistency=1, |
|
step_size=1, |
|
max_iter=150, |
|
tol=1e-4, |
|
cg_max_iter=5000, |
|
cg_tol=1e-3, |
|
cut_intersection=True, |
|
): |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
num_normals = cp.sum(normal_mask) |
|
normal_map_front = cp.asarray(normal_front) |
|
normal_map_back = cp.asarray(normal_back) |
|
normal_mask = cp.asarray(normal_mask) |
|
if depth_mask is not None: |
|
depth_map_front = cp.asarray(depth_front) |
|
depth_map_back = cp.asarray(depth_back) |
|
depth_mask = cp.asarray(depth_mask) |
|
|
|
|
|
nx_front = normal_map_front[normal_mask, 1] |
|
ny_front = normal_map_front[normal_mask, 0] |
|
nz_front = -normal_map_front[normal_mask, 2] |
|
del normal_map_front |
|
|
|
nx_back = normal_map_back[normal_mask, 1] |
|
ny_back = normal_map_back[normal_mask, 0] |
|
nz_back = -normal_map_back[normal_mask, 2] |
|
del normal_map_back |
|
|
|
|
|
A3_f, A4_f, A1_f, A2_f = generate_dx_dy( |
|
normal_mask, nz_horizontal=nz_front, nz_vertical=nz_front, step_size=step_size |
|
) |
|
A3_b, A4_b, A1_b, A2_b = generate_dx_dy( |
|
normal_mask, nz_horizontal=nz_back, nz_vertical=nz_back, step_size=step_size |
|
) |
|
|
|
has_left_mask = cp.logical_and(move_right(normal_mask), normal_mask) |
|
has_right_mask = cp.logical_and(move_left(normal_mask), normal_mask) |
|
has_bottom_mask = cp.logical_and(move_top(normal_mask), normal_mask) |
|
has_top_mask = cp.logical_and(move_bottom(normal_mask), normal_mask) |
|
|
|
top_boundnary_mask = cp.logical_xor(has_top_mask, normal_mask)[normal_mask] |
|
bottom_boundary_mask = cp.logical_xor(has_bottom_mask, normal_mask)[normal_mask] |
|
left_boundary_mask = cp.logical_xor(has_left_mask, normal_mask)[normal_mask] |
|
right_boudnary_mask = cp.logical_xor(has_right_mask, normal_mask)[normal_mask] |
|
|
|
A_front_data = vstack((A1_f, A2_f, A3_f, A4_f)) |
|
A_front_zero = csr_matrix(A_front_data.shape) |
|
A_front = hstack([A_front_data, A_front_zero]) |
|
|
|
A_back_data = vstack((A1_b, A2_b, A3_b, A4_b)) |
|
A_back_zero = csr_matrix(A_back_data.shape) |
|
A_back = hstack([A_back_zero, A_back_data]) |
|
|
|
b_front = cp.concatenate((-nx_front, -nx_front, -ny_front, -ny_front)) |
|
b_back = cp.concatenate((-nx_back, -nx_back, -ny_back, -ny_back)) |
|
|
|
|
|
W_front = spdiags( |
|
0.5 * cp.ones(4 * num_normals), 0, 4 * num_normals, 4 * num_normals, format="csr" |
|
) |
|
W_back = spdiags( |
|
0.5 * cp.ones(4 * num_normals), 0, 4 * num_normals, 4 * num_normals, format="csr" |
|
) |
|
|
|
z_front = cp.zeros(num_normals, float) |
|
z_back = cp.zeros(num_normals, float) |
|
z_combined = cp.concatenate((z_front, z_back)) |
|
|
|
B, B_full = create_boundary_matrix(normal_mask) |
|
B_mat = lambda_boundary_consistency * coo_matrix(B_full.get().T @ B_full.get()) |
|
|
|
energy_list = [] |
|
|
|
if depth_mask is not None: |
|
depth_mask_flat = depth_mask[normal_mask].astype(bool) |
|
z_prior_front = depth_map_front[normal_mask] |
|
z_prior_front[~depth_mask_flat] = 0 |
|
z_prior_back = depth_map_back[normal_mask] |
|
z_prior_back[~depth_mask_flat] = 0 |
|
m = depth_mask[normal_mask].astype(int) |
|
M = diags(m) |
|
|
|
energy = (A_front @ z_combined - b_front).T @ W_front @ (A_front @ z_combined - b_front) + \ |
|
lambda_normal_back * (A_back @ z_combined - b_back).T @ W_back @ (A_back @ z_combined - b_back) + \ |
|
lambda_depth_front * (z_front - z_prior_front).T @ M @ (z_front - z_prior_front) + \ |
|
lambda_depth_back * (z_back - z_prior_back).T @ M @ (z_back - z_prior_back) + \ |
|
lambda_boundary_consistency * (z_back - z_front).T @ B @ (z_back - z_front) |
|
|
|
depth_map_front_est = cp.ones_like(normal_mask, float) * cp.nan |
|
depth_map_back_est = cp.ones_like(normal_mask, float) * cp.nan |
|
|
|
facets_back = cp.asnumpy(construct_facets_from(normal_mask)) |
|
faces_back = np.concatenate((facets_back[:, [1, 4, 3]], facets_back[:, [1, 3, 2]]), axis=0) |
|
faces_front = np.concatenate((facets_back[:, [1, 2, 3]], facets_back[:, [1, 3, 4]]), axis=0) |
|
|
|
for i in range(max_iter): |
|
A_mat_front = A_front_data.T @ W_front @ A_front_data |
|
b_vec_front = A_front_data.T @ W_front @ b_front |
|
|
|
A_mat_back = A_back_data.T @ W_back @ A_back_data |
|
b_vec_back = A_back_data.T @ W_back @ b_back |
|
if depth_mask is not None: |
|
b_vec_front += lambda_depth_front * M @ z_prior_front |
|
b_vec_back += lambda_depth_back * M @ z_prior_back |
|
A_mat_front += lambda_depth_front * M |
|
A_mat_back += lambda_depth_back * M |
|
offset_front = cp.mean((z_prior_front - z_combined[:num_normals])[depth_mask_flat]) |
|
offset_back = cp.mean((z_prior_back - z_combined[num_normals:])[depth_mask_flat]) |
|
z_combined[:num_normals] = z_combined[:num_normals] + offset_front |
|
z_combined[num_normals:] = z_combined[num_normals:] + offset_back |
|
|
|
|
|
A_mat_combined = hstack([vstack((A_mat_front, csr_matrix((num_normals, num_normals)))), \ |
|
vstack((csr_matrix((num_normals, num_normals)), A_mat_back))]) + B_mat |
|
b_vec_combined = cp.concatenate((b_vec_front, b_vec_back)) |
|
|
|
D = spdiags( |
|
1 / cp.clip(A_mat_combined.diagonal(), 1e-5, None), 0, 2 * num_normals, 2 * num_normals, |
|
"csr" |
|
) |
|
|
|
z_combined, _ = cg( |
|
A_mat_combined, b_vec_combined, M=D, x0=z_combined, maxiter=cg_max_iter, tol=cg_tol |
|
) |
|
z_front = z_combined[:num_normals] |
|
z_back = z_combined[num_normals:] |
|
wu_f = sigmoid((A2_f.dot(z_front))**2 - (A1_f.dot(z_front))**2, k) |
|
wv_f = sigmoid((A4_f.dot(z_front))**2 - (A3_f.dot(z_front))**2, k) |
|
wu_f[top_boundnary_mask] = 0.5 |
|
wu_f[bottom_boundary_mask] = 0.5 |
|
wv_f[left_boundary_mask] = 0.5 |
|
wv_f[right_boudnary_mask] = 0.5 |
|
W_front = spdiags( |
|
cp.concatenate((wu_f, 1 - wu_f, wv_f, 1 - wv_f)), |
|
0, |
|
4 * num_normals, |
|
4 * num_normals, |
|
format="csr" |
|
) |
|
|
|
wu_b = sigmoid((A2_b.dot(z_back))**2 - (A1_b.dot(z_back))**2, k) |
|
wv_b = sigmoid((A4_b.dot(z_back))**2 - (A3_b.dot(z_back))**2, k) |
|
wu_b[top_boundnary_mask] = 0.5 |
|
wu_b[bottom_boundary_mask] = 0.5 |
|
wv_b[left_boundary_mask] = 0.5 |
|
wv_b[right_boudnary_mask] = 0.5 |
|
W_back = spdiags( |
|
cp.concatenate((wu_b, 1 - wu_b, wv_b, 1 - wv_b)), |
|
0, |
|
4 * num_normals, |
|
4 * num_normals, |
|
format="csr" |
|
) |
|
|
|
energy_old = energy |
|
energy = (A_front_data @ z_front - b_front).T @ W_front @ (A_front_data @ z_front - b_front) + \ |
|
lambda_normal_back * (A_back_data @ z_back - b_back).T @ W_back @ (A_back_data @ z_back - b_back) + \ |
|
lambda_depth_front * (z_front - z_prior_front).T @ M @ (z_front - z_prior_front) + \ |
|
lambda_depth_back * (z_back - z_prior_back).T @ M @ (z_back - z_prior_back) +\ |
|
lambda_boundary_consistency * (z_back - z_front).T @ B @ (z_back - z_front) |
|
|
|
energy_list.append(energy) |
|
relative_energy = cp.abs(energy - energy_old) / energy_old |
|
|
|
|
|
|
|
|
|
if False: |
|
|
|
depth_map_front_est[normal_mask] = z_front |
|
depth_map_back_est[normal_mask] = z_back |
|
vertices_front = cp.asnumpy( |
|
map_depth_map_to_point_clouds( |
|
depth_map_front_est, normal_mask, K=None, step_size=step_size |
|
) |
|
) |
|
vertices_back = cp.asnumpy( |
|
map_depth_map_to_point_clouds( |
|
depth_map_back_est, normal_mask, K=None, step_size=step_size |
|
) |
|
) |
|
|
|
vertices_front, faces_front_ = remove_stretched_faces(vertices_front, faces_front) |
|
vertices_back, faces_back_ = remove_stretched_faces(vertices_back, faces_back) |
|
|
|
F_verts = verts_inverse_transform(torch.as_tensor(vertices_front).float(), 256.0) |
|
B_verts = verts_inverse_transform(torch.as_tensor(vertices_back).float(), 256.0) |
|
|
|
F_B_verts = torch.cat((F_verts, B_verts), dim=0) |
|
F_B_faces = torch.cat(( |
|
torch.as_tensor(faces_front_).long(), |
|
torch.as_tensor(faces_back_).long() + faces_front_.max() + 1 |
|
), |
|
dim=0) |
|
|
|
front_surf = trimesh.Trimesh(F_verts, faces_front_) |
|
back_surf = trimesh.Trimesh(B_verts, faces_back_) |
|
double_surf = trimesh.Trimesh(F_B_verts, F_B_faces) |
|
|
|
bini_dir = "/home/yxiu/Code/ECON/log/bini/OBJ" |
|
front_surf.export(osp.join(bini_dir, f"{i:04d}_F.obj")) |
|
back_surf.export(osp.join(bini_dir, f"{i:04d}_B.obj")) |
|
double_surf.export(osp.join(bini_dir, f"{i:04d}_FB.obj")) |
|
|
|
if relative_energy < tol: |
|
break |
|
|
|
|
|
depth_map_front_est[normal_mask] = z_front |
|
depth_map_back_est[normal_mask] = z_back |
|
|
|
if cut_intersection: |
|
|
|
normal_mask[depth_map_front_est >= depth_map_back_est] = False |
|
depth_map_front_est[~normal_mask] = cp.nan |
|
depth_map_back_est[~normal_mask] = cp.nan |
|
|
|
vertices_front = cp.asnumpy( |
|
map_depth_map_to_point_clouds( |
|
depth_map_front_est, normal_mask, K=None, step_size=step_size |
|
) |
|
) |
|
vertices_back = cp.asnumpy( |
|
map_depth_map_to_point_clouds(depth_map_back_est, normal_mask, K=None, step_size=step_size) |
|
) |
|
|
|
facets_back = cp.asnumpy(construct_facets_from(normal_mask)) |
|
faces_back = np.concatenate((facets_back[:, [1, 4, 3]], facets_back[:, [1, 3, 2]]), axis=0) |
|
faces_front = np.concatenate((facets_back[:, [1, 2, 3]], facets_back[:, [1, 3, 4]]), axis=0) |
|
|
|
vertices_front, faces_front = remove_stretched_faces(vertices_front, faces_front) |
|
vertices_back, faces_back = remove_stretched_faces(vertices_back, faces_back) |
|
|
|
front_mesh = clean_floats(trimesh.Trimesh(vertices_front, faces_front)) |
|
back_mesh = clean_floats(trimesh.Trimesh(vertices_back, faces_back)) |
|
|
|
result = { |
|
"F_verts": torch.as_tensor(front_mesh.vertices).float(), "F_faces": torch.as_tensor( |
|
front_mesh.faces |
|
).long(), "B_verts": torch.as_tensor(back_mesh.vertices).float(), "B_faces": |
|
torch.as_tensor(back_mesh.faces).long(), "F_depth": |
|
torch.as_tensor(depth_map_front_est).float(), "B_depth": |
|
torch.as_tensor(depth_map_back_est).float() |
|
} |
|
|
|
return result |
|
|
|
|
|
def save_normal_tensor(in_tensor, idx, png_path, thickness=0.0): |
|
|
|
os.makedirs(os.path.dirname(png_path), exist_ok=True) |
|
|
|
normal_F_arr = tensor2arr(in_tensor["normal_F"][idx:idx + 1]) |
|
normal_B_arr = tensor2arr(in_tensor["normal_B"][idx:idx + 1]) |
|
mask_normal_arr = tensor2arr(in_tensor["image"][idx:idx + 1], True) |
|
|
|
depth_F_arr = depth2arr(in_tensor["depth_F"][idx]) |
|
depth_B_arr = depth2arr(in_tensor["depth_B"][idx]) |
|
|
|
BNI_dict = {} |
|
|
|
|
|
BNI_dict["normal_F"] = normal_F_arr |
|
BNI_dict["normal_B"] = normal_B_arr |
|
BNI_dict["mask"] = mask_normal_arr > 0. |
|
BNI_dict["depth_F"] = depth_F_arr - 100. - thickness |
|
BNI_dict["depth_B"] = 100. - depth_B_arr + thickness |
|
BNI_dict["depth_mask"] = depth_F_arr != -1.0 |
|
|
|
np.save(png_path + ".npy", BNI_dict, allow_pickle=True) |
|
|
|
return BNI_dict |
|
|