Spaces:
Runtime error
Runtime error
| 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 | |