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) # green
cv2.drawContours(disp, contour2, -1, (255, 0, 0), 1) # blue
for i in range(contour1.shape[1]): # do not show all the points when display
# cv2.circle(disp, (contour1[0, i, 0, 0], contour1[0, i, 0, 1]), 1,
# (255, 0, 0), -1)
corresPoint = contour2[0, phi[i], 0]
# cv2.circle(disp, (corresPoint[0], corresPoint[1]), 1, (0, 255, 0), -1)
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)
# cos(90-20)=0.34 cos(90-10)=0.17, 10~20 degree
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
# BNI related
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):
# pixel coordinates
# ^ vertical positive
# |
# |
# |
# o ---> horizontal positive
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):
# pixel coordinates
# ^ vertical positive
# |
# |
# |
# o ---> horizontal positive
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):
# y
# | z
# | /
# |/
# o ---x
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 # 3 x m
vertices = (cp.linalg.inv(K) @ u).T * depth_map[mask, cp.newaxis] # m x 3
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
# discard single point
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,
):
# To avoid confusion, we list the coordinate systems in this code as follows
#
# pixel coordinates camera coordinates normal coordinates (the main paper's Fig. 1 (a))
# u x y
# | | z |
# | | / o -- x
# | |/ /
# o --- v o --- y z
# (bottom left)
# (o is the optical center;
# xy-plane is parallel to the image plane;
# +z is the viewing direction.)
#
# The input normal map should be defined in the normal coordinates.
# The camera matrix K should be defined in the camera coordinates.
# K = [[fx, 0, cx],
# [0, fy, cy],
# [0, 0, 1]]
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)
# transfer the normal map from the normal coordinates to the camera coordinates
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
# right, left, top, bottom
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))
# initialization
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()) #bug
energy_list = []
if depth_mask is not None:
depth_mask_flat = depth_mask[normal_mask].astype(bool) # shape: (num_normals,)
z_prior_front = depth_map_front[normal_mask] # shape: (num_normals,)
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"
) # Jacob preconditioner
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) # top
wv_f = sigmoid((A4_f.dot(z_front))**2 - (A3_f.dot(z_front))**2, k) # right
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) # top
wv_b = sigmoid((A4_b.dot(z_back))**2 - (A3_b.dot(z_back))**2, k) # right
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
# print(f"step {i + 1}/{max_iter} energy: {energy:.3e}"
# f" relative energy: {relative_energy:.3e}")
if False:
# intermediate results
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
# del A1, A2, A3, A4, nx, ny
depth_map_front_est[normal_mask] = z_front
depth_map_back_est[normal_mask] = z_back
if cut_intersection:
# manually cut the 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 = {}
# clothed human
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