main.py

# Copyright (c) Meta Platforms, Inc. and affiliates.

import argparse
import datetime
import numpy as np
import os
import pyvista as pv
import time
import trimesh
import yaml

from scipy.sparse import csr_matrix, spdiags
from scipy.sparse.linalg import cg

from bini.bilateral_normal_integration import (
    bilateral_normal_integration,
    construct_facets_from,
)
from bini.conversions import (
    build_gamma_equation_matrix_from_bini,
    compute_residual_and_wuv_image,
)
from datasets.diligent import DiLiGenTObject
from datasets.diligent_mv import DiLiGenTMVObject
from utils.metrics import MADEComputer
from utils.parsing import float_or_string, float_or_none


def compute_b_ours_log_with_gamma_A(
    A_ours_log_with_gamma,
    indices_z_a,
    H,
    W,
    num_shifts,
    v_where_is_valid_and_valid_neighbor,
    u_where_is_valid_and_valid_neighbor,
    shift_where_is_valid_and_valid_neighbor,
    residual_image,
    wuv_image,
    mask,
    w_b_to_a,
    w_eps_a,
    rho_beta,
    q_beta,
):
    assert (rho_beta is None) == (q_beta is None)
    gamma_ = np.array(
        A_ours_log_with_gamma[
            np.arange(len(indices_z_a)),
            indices_z_a,
        ]
    )[0]
    if rho_beta is None:
        # No discontinuity computation.
        alpha_times_beta = np.zeros_like(wuv_image)
    else:
        # Compute alpha's.
        gamma_image = np.zeros((H, W, num_shifts))
        gamma_image[:] = np.nan
        gamma_image[
            v_where_is_valid_and_valid_neighbor,
            u_where_is_valid_and_valid_neighbor,
            shift_where_is_valid_and_valid_neighbor,
        ] = gamma_
        alpha = residual_image.copy()
        alpha[np.logical_not(mask)] = np.nan
        # alpha is eps_a_from_b / z_b.
        alpha = (np.exp(alpha / gamma_image) - w_b_to_a) / w_eps_a[..., None]
        # We want beta to be roughly 1 when w_1 - w_2 <= -1 and roughly 0 when
        # w_1 - w_2 > -1,
        # This can be achieved for instance with a sigmoid S(x) reflected around the y
        # axis and centered at -1 + eps, where eps > 0 is a small quantity.
        beta_image = 1.0 / (1 + np.exp(q_beta * (wuv_image - rho_beta)))
        alpha_times_beta = alpha * beta_image
    # Update b.
    b_ours_log_with_gamma_A = (
        np.log(w_b_to_a + alpha_times_beta * w_eps_a[..., None])[
            v_where_is_valid_and_valid_neighbor,
            u_where_is_valid_and_valid_neighbor,
            shift_where_is_valid_and_valid_neighbor,
        ]
    ) * gamma_
    b_ours_log_with_gamma_A[np.isnan(b_ours_log_with_gamma_A)] = 0.0

    return b_ours_log_with_gamma_A


def perform_optimization(
    A_ours_log_with_gamma,
    H,
    W,
    num_shifts,
    v_where_is_valid_and_valid_neighbor,
    u_where_is_valid_and_valid_neighbor,
    shift_where_is_valid_and_valid_neighbor,
    indices_z_a,
    mask,
    w_b_to_a,
    w_eps_a,
    rho_beta,
    q_beta,
    force_all_iters=False,
    sigmoid_k_value=2,
    max_iter=150,
    cg_max_iter=5000,
    cg_tol=1.0e-3,
    tol=1.0e-4,
):
    num_equations, num_is_valid_pixel = A_ours_log_with_gamma.shape

    W_matrix = spdiags(
        0.5 * np.ones(num_equations), 0, num_equations, num_equations, format="csr"
    )

    log_z = np.zeros(num_is_valid_pixel)

    residual_image, wuv_image = compute_residual_and_wuv_image(
        A_ours_log_with_gamma=A_ours_log_with_gamma,
        log_z=log_z,
        H=H,
        W=W,
        num_shifts=num_shifts,
        v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor=(
            shift_where_is_valid_and_valid_neighbor
        ),
        sigmoid_k_value=sigmoid_k_value,
    )

    # Compute initial b term assuming no discontinuities.
    b_ours_log_with_gamma_A = compute_b_ours_log_with_gamma_A(
        A_ours_log_with_gamma=A_ours_log_with_gamma,
        indices_z_a=indices_z_a,
        H=H,
        W=W,
        num_shifts=num_shifts,
        v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor=(
            shift_where_is_valid_and_valid_neighbor
        ),
        residual_image=residual_image,
        wuv_image=wuv_image,
        mask=mask,
        w_b_to_a=w_b_to_a,
        w_eps_a=w_eps_a,
        rho_beta=None,
        q_beta=None,
    )

    energy = (
        (A_ours_log_with_gamma @ log_z - b_ours_log_with_gamma_A).T
        @ W_matrix
        @ (A_ours_log_with_gamma @ log_z - b_ours_log_with_gamma_A)
    )

    tic = time.time()

    log_z_list = []
    made_list = []
    energy_list = []
    pbar = range(max_iter)

    # Optimization loop
    log_z_list.append(log_z)
    try:
        log_z_image = np.ones((H, W)) * np.nan
        log_z_image[is_valid_pixel_mask] = log_z
        made_list.append(made_computer.compute_curr_made(log_z=log_z_image)[0])
    except ValueError:
        pass

    for i in pbar:
        # fix weights and solve for depths
        A_mat = A_ours_log_with_gamma.T @ W_matrix @ A_ours_log_with_gamma
        b_vec = A_ours_log_with_gamma.T @ W_matrix @ b_ours_log_with_gamma_A
        D = spdiags(
            1 / np.clip(A_mat.diagonal(), 1e-5, None),
            0,
            num_is_valid_pixel,
            num_is_valid_pixel,
            format="csr",
        )  # Jacob preconditioner

        log_z, _ = cg(A_mat, b_vec, x0=log_z, M=D, maxiter=cg_max_iter, tol=cg_tol)
        log_z_list.append(log_z)
        try:
            log_z_image = np.ones((H, W)) * np.nan
            log_z_image[is_valid_pixel_mask] = log_z
            made_list.append(made_computer.compute_curr_made(log_z=log_z_image)[0])
        except ValueError:
            pass

        # Update the weight matrices.
        residual_image, wuv_image = compute_residual_and_wuv_image(
            A_ours_log_with_gamma=A_ours_log_with_gamma,
            log_z=log_z,
            H=H,
            W=W,
            num_shifts=num_shifts,
            v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
            u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
            shift_where_is_valid_and_valid_neighbor=(
                shift_where_is_valid_and_valid_neighbor
            ),
            sigmoid_k_value=sigmoid_k_value,
        )

        W_matrix = wuv_image[
            v_where_is_valid_and_valid_neighbor,
            u_where_is_valid_and_valid_neighbor,
            shift_where_is_valid_and_valid_neighbor,
        ]
        W_matrix = spdiags(W_matrix, 0, num_equations, num_equations, format="csr")

        # Check for convergence.
        energy_old = energy
        energy = (
            (A_ours_log_with_gamma @ log_z - b_ours_log_with_gamma_A).T
            @ W_matrix
            @ (A_ours_log_with_gamma @ log_z - b_ours_log_with_gamma_A)
        )
        energy_list.append(energy)
        relative_energy = np.abs(energy - energy_old) / energy_old

        if relative_energy < tol and (not force_all_iters):
            break

        b_ours_log_with_gamma_A = compute_b_ours_log_with_gamma_A(
            A_ours_log_with_gamma=A_ours_log_with_gamma,
            indices_z_a=indices_z_a,
            H=H,
            W=W,
            num_shifts=num_shifts,
            v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
            u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
            shift_where_is_valid_and_valid_neighbor=(
                shift_where_is_valid_and_valid_neighbor
            ),
            residual_image=residual_image,
            wuv_image=wuv_image,
            mask=mask,
            w_b_to_a=w_b_to_a,
            w_eps_a=w_eps_a,
            rho_beta=rho_beta,
            q_beta=q_beta,
        )
    toc = time.time()

    print(f"Total time: {toc - tic:.3f} sec")

    return (made_list, log_z)


parser = argparse.ArgumentParser()

# Examples: ".../DiLiGenT", ".../DiLiGenT-MV".
parser.add_argument("--data-dir", required=True, type=str)
parser.add_argument(
    "--dataset-type", required=True, choices=["diligent", "diligent_mv"]
)
parser.add_argument("--obj-name", required=True, type=str)
parser.add_argument(
    "--num-shifts", type=str, choices=["4", "8", "4_diag"], required=True
)
parser.add_argument("--lambda-m", required=True, type=float_or_string)
parser.add_argument("--k-sigmoid-lambda-m", type=float)
parser.add_argument("--normal-type", type=str, default="gt")
parser.add_argument("--max-iter", type=int, default=150)
parser.add_argument("--rho-beta", required=True, type=float_or_none)
parser.add_argument("--q-beta", required=True, type=float_or_none)
parser.add_argument("--gamma-type", required=True, choices=["ours", "bini"])
parser.add_argument("--w-b-to-a-outlier-th", type=float, required=True)  # 1.1
parser.add_argument("--num-ms", type=int, required=True)  # 15
parser.add_argument("--bad-normals-correction-criterion", type=str, default="NaN")
parser.add_argument("--threshold-rel-change-abs-n-dot-tau", type=float_or_none)
parser.add_argument("--force-all-iters", action="store_true")
parser.add_argument("--output-folder", type=str, required=True)

args = parser.parse_args()
data_dir = args.data_dir
dataset_type = args.dataset_type
obj_name = args.obj_name
num_shifts = args.num_shifts
lambda_m = args.lambda_m
k_sigmoid_lambda_m = args.k_sigmoid_lambda_m
normal_type = args.normal_type
max_iter = args.max_iter
rho_beta = args.rho_beta
q_beta = args.q_beta
gamma_type = args.gamma_type
bad_normals_correction_criterion = args.bad_normals_correction_criterion
threshold_rel_change_abs_n_dot_tau = args.threshold_rel_change_abs_n_dot_tau
force_all_iters = args.force_all_iters
output_folder = args.output_folder
w_b_to_a_outlier_th = args.w_b_to_a_outlier_th
num_ms = args.num_ms

if dataset_type == "diligent":
    object_class = DiLiGenTObject
elif dataset_type == "diligent_mv":
    object_class = DiLiGenTMVObject

diligent_object = object_class(
    data_dir=data_dir,
    obj_name=obj_name,
    normal_type=normal_type,
    num_shifts=num_shifts,
    lambda_m=lambda_m,
    k_sigmoid_lambda_m=k_sigmoid_lambda_m,
    w_b_to_a_outlier_th=w_b_to_a_outlier_th,
    num_ms=num_ms,
    threshold_rel_change_abs_n_dot_tau=threshold_rel_change_abs_n_dot_tau,
    bad_normals_correction_criterion=bad_normals_correction_criterion,
)

# Set up output folders, formatting their filename using date and time.
common_subfolder_str = (
    datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + f"_{diligent_object.obj_name}"
)
if dataset_type == "diligent_mv":
    output_folder = os.path.join(output_folder, normal_type)
    common_subfolder_str = common_subfolder_str + f"_view_{diligent_object.view_idx}"
output_folder = os.path.join(output_folder, common_subfolder_str)
os.makedirs(output_folder)

is_valid_pixel_mask = diligent_object.is_valid_pixel_mask
depth_gt = diligent_object.depth_gt
K = diligent_object.K
n_a_vec = diligent_object.n_a_vec
tau_a_vec = diligent_object.tau_a_vec
w_b_to_a = diligent_object.w_b_to_a
w_eps_a = diligent_object.w_eps_a

v_where_is_valid = diligent_object.v_where_is_valid
u_where_is_valid = diligent_object.u_where_is_valid
v_where_is_valid_and_valid_neighbor = (
    diligent_object.v_where_is_valid_and_valid_neighbor
)
u_where_is_valid_and_valid_neighbor = (
    diligent_object.u_where_is_valid_and_valid_neighbor
)
shift_where_is_valid_and_valid_neighbor = (
    diligent_object.shift_where_is_valid_and_valid_neighbor
)
H = diligent_object.H
W = diligent_object.W
num_shifts = diligent_object.num_shifts
channel_idx_to_du_dv = diligent_object.channel_idx_to_du_dv

made_computer = MADEComputer(
    is_valid_pixel_mask=is_valid_pixel_mask, log_depth_gt=np.log(depth_gt)
)

# Reconstruct with BiNI.
(
    depth_map_est_bini,
    nz_u_bini,
    nz_v_bini,
    mask_bini,
) = bilateral_normal_integration(
    normal_map=n_a_vec,
    normal_mask=is_valid_pixel_mask,
    k=2,
    K=K,
    max_iter=max_iter,
    tol=1e-4,
    force_all_iters=force_all_iters,
)

made_bini, _, scale_bini = made_computer.compute_curr_made(
    log_z=np.log(depth_map_est_bini), return_scale=True
)
print(f"\033[1mMADE BiNI = {made_bini}\033[0m")

# Reconstruct with ours.
should_compute_bini_coeffs = num_shifts == 4
(
    A_bini,
    (indices_z_a, indices_z_b),
) = build_gamma_equation_matrix_from_bini(
    v_where_is_valid=v_where_is_valid,
    u_where_is_valid=u_where_is_valid,
    v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
    u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
    shift_where_is_valid_and_valid_neighbor=shift_where_is_valid_and_valid_neighbor,
    H=H,
    W=W,
    num_shifts=num_shifts,
    channel_idx_to_du_dv=channel_idx_to_du_dv,
    nz_u_bini=nz_u_bini if should_compute_bini_coeffs else None,
    nz_v_bini=nz_v_bini if should_compute_bini_coeffs else None,
    mask_bini=mask_bini if should_compute_bini_coeffs else None,
    normal_image_ours_convention=n_a_vec.copy() if should_compute_bini_coeffs else None,
)

# - Define the multiplicative factor for each equation, either using the original BiNI
#   formulation, or using our generic formulation.
if gamma_type == "bini":
    A_ours_log_with_gamma = A_bini.copy()
elif gamma_type == "ours":
    A_ours_log_with_gamma = csr_matrix(A_bini.shape)
    curr_duv_ = channel_idx_to_du_dv[shift_where_is_valid_and_valid_neighbor]
    curr_du_ = curr_duv_[..., 0]
    curr_dv_ = curr_duv_[..., 1]
    # gamma_ = ||u_b - u_a|| / ||tau_b - tau_a|| * n_a.tau_a.
    gamma_ = (
        np.linalg.norm(
            channel_idx_to_du_dv[shift_where_is_valid_and_valid_neighbor], axis=-1
        )
        / np.linalg.norm(
            tau_a_vec[
                v_where_is_valid_and_valid_neighbor + curr_dv_,
                u_where_is_valid_and_valid_neighbor + curr_du_,
            ]
            - tau_a_vec[
                v_where_is_valid_and_valid_neighbor,
                u_where_is_valid_and_valid_neighbor,
            ],
            axis=-1,
        )
    ) * np.einsum("ijk,ijk->ij", n_a_vec, tau_a_vec)[
        v_where_is_valid_and_valid_neighbor, u_where_is_valid_and_valid_neighbor
    ]

    # z_a entries.
    A_ours_log_with_gamma[np.arange(len(indices_z_a)), indices_z_a] = gamma_
    # z_b entries.
    A_ours_log_with_gamma[np.arange(len(indices_z_b)), indices_z_b] = -gamma_
else:
    raise ValueError(f"Invalid gamma type '{gamma_type}'.")


made_list = []

(curr_made_list, log_z) = perform_optimization(
    A_ours_log_with_gamma=A_ours_log_with_gamma,
    H=H,
    W=W,
    num_shifts=num_shifts,
    v_where_is_valid_and_valid_neighbor=v_where_is_valid_and_valid_neighbor,
    u_where_is_valid_and_valid_neighbor=u_where_is_valid_and_valid_neighbor,
    shift_where_is_valid_and_valid_neighbor=shift_where_is_valid_and_valid_neighbor,
    indices_z_a=indices_z_a,
    mask=is_valid_pixel_mask,
    w_b_to_a=w_b_to_a,
    w_eps_a=w_eps_a,
    rho_beta=rho_beta,
    q_beta=q_beta,
    force_all_iters=force_all_iters,
    sigmoid_k_value=2,
    max_iter=max_iter,
    cg_max_iter=5000,
    cg_tol=1.0e-3,
    tol=1.0e-4,
)
made_list += curr_made_list

if len(made_list) == 0:
    print("\033[91mNo MADEs were computed.\033[0m")
else:
    np.savetxt(os.path.join(output_folder, "made_list.txt"), np.array(made_list))

# - Reconstruct the depth map and surface
depth_map = np.zeros((H, W)) * np.nan
depth_map[v_where_is_valid, u_where_is_valid] = log_z
depth_map = np.exp(depth_map)

made_ours_log_with_gamma_A, _, scale_ours = made_computer.compute_curr_made(
    log_z=np.log(depth_map.copy()), return_scale=True
)

print(f"\033[1mMADE Ours = {made_ours_log_with_gamma_A}\033[0m")

# - Save point cloud.
reconstructed_point_image = {}
surface = {}

facets = construct_facets_from(is_valid_pixel_mask)
if n_a_vec.copy()[:, :, -1].mean() < 0:
    facets = facets[:, [0, 1, 4, 3, 2]]

normal_image_vis = n_a_vec.copy()
#   - Revert coordinate frame for visualization.
normal_image_vis[..., 1:] *= -1
normal_image_vis = (1.0 + normal_image_vis) / 2.0

for curr_method_name, curr_depth_map in [
    ("bini", depth_map_est_bini * scale_bini),
    ("ours", depth_map * scale_ours),
    ("gt", depth_gt),
]:
    if curr_depth_map is None:
        continue
    reconstructed_point_image[curr_method_name] = np.zeros((H, W, 3))
    reconstructed_point_image[curr_method_name][:] = np.nan
    reconstructed_point_image[curr_method_name][
        v_where_is_valid, u_where_is_valid, 2
    ] = curr_depth_map[v_where_is_valid, u_where_is_valid]
    reconstructed_point_image[curr_method_name][
        v_where_is_valid, u_where_is_valid, 0
    ] = (
        tau_a_vec[v_where_is_valid, u_where_is_valid, 0]
        * curr_depth_map[v_where_is_valid, u_where_is_valid]
    )
    reconstructed_point_image[curr_method_name][
        v_where_is_valid, u_where_is_valid, 1
    ] = (
        tau_a_vec[v_where_is_valid, u_where_is_valid, 1]
        * curr_depth_map[v_where_is_valid, u_where_is_valid]
    )
    valid_points = reconstructed_point_image[curr_method_name][
        v_where_is_valid, u_where_is_valid
    ]
    surface[curr_method_name] = pv.PolyData(valid_points, facets)
    surface[curr_method_name].save(
        os.path.join(output_folder, f"{obj_name}_{curr_method_name}_mesh.ply"),
        binary=False,
    )

    pc = trimesh.PointCloud(
        vertices=valid_points,
        colors=normal_image_vis[v_where_is_valid, u_where_is_valid],
    )
    output_path = os.path.join(output_folder, f"{obj_name}_{curr_method_name}_pc.ply")
    print(output_path)
    _ = pc.export(output_path)

# - Save parameters and MADEs to a YAML file.
with open(os.path.join(output_folder, "config.yaml"), "w") as f:
    yaml.dump(
        {
            **vars(args),
            "made_ours": float(made_ours_log_with_gamma_A),
            "made_bini": float(made_bini),
        },
        f,
    )