lib/common/BNI_utils.py

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