conversions.py

"""
Copyright (c) 2022 Xu Cao
Copyright (c) Meta Platforms, Inc. and affiliates.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

Partly adapted from https://github.com/xucao-42/bilateral_normal_integration/.
"""

import numpy as np
import scipy

from bini.bilateral_normal_integration import (
    move_bottom,
    move_left,
    move_right,
    move_top,
)
from bini.misc import sigmoid


def build_gamma_equation_matrix_from_bini(
    v_where_is_valid: np.ndarray,
    u_where_is_valid: np.ndarray,
    v_where_is_valid_and_valid_neighbor: np.ndarray,
    u_where_is_valid_and_valid_neighbor: np.ndarray,
    shift_where_is_valid_and_valid_neighbor: np.ndarray,
    H: int,
    W: int,
    num_shifts: int,
    channel_idx_to_du_dv: np.ndarray,
    nz_u_bini: np.ndarray,
    nz_v_bini: np.ndarray,
    mask_bini: np.ndarray,
    normal_image_ours_convention: np.ndarray,
):
    assert num_shifts == 4
    direction_to_shift = {"left": 0, "top": 1, "right": 2, "bottom": 3}

    assert nz_u_bini.shape == nz_v_bini.shape == (len(v_where_is_valid),)
    assert mask_bini.shape == (H, W)
    assert normal_image_ours_convention.shape == (H, W, 3)
    mask_own = np.zeros((H, W), dtype=bool)
    mask_own[v_where_is_valid, u_where_is_valid] = True
    assert np.all(mask_bini == mask_own)
    nz_u_image = np.zeros((H, W))
    nz_u_image[:] = np.nan
    nz_v_image = nz_u_image.copy()
    nz_u_image[v_where_is_valid, u_where_is_valid] = nz_u_bini
    nz_v_image[v_where_is_valid, u_where_is_valid] = nz_v_bini
    has_mask = {
        "left": np.logical_and(move_right(mask_bini), mask_bini),
        "right": np.logical_and(move_left(mask_bini), mask_bini),
        "bottom": np.logical_and(move_top(mask_bini), mask_bini),
        "top": np.logical_and(move_bottom(mask_bini), mask_bini),
    }

    neighbor_masks_own = np.zeros((H, W, num_shifts), dtype=bool)
    neighbor_masks_own[
        v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor,
    ] = True

    for direction, shift in direction_to_shift.items():
        assert np.all(neighbor_masks_own[..., shift] == has_mask[direction])
    nz_bini = {}
    nxy_bini = {}
    nz_bini_image = np.zeros((H, W, num_shifts))
    nz_bini_image[:] = np.nan
    nxy_bini_image = nz_bini_image.copy()
    for direction, shift in direction_to_shift.items():
        if direction in ["left", "right"]:
            curr_nz_uv = -nz_u_image if direction == "left" else nz_u_image
            curr_nxy_uv = normal_image_ours_convention[..., 0]
        elif direction in ["top", "bottom"]:
            # NOTE: This is a necessary step to convert normals from our convention to
            # the convention used by BiNI.
            curr_nz_uv = -nz_v_image if direction == "bottom" else nz_v_image
            curr_nxy_uv = -normal_image_ours_convention[..., 1]
        nz_bini[direction] = curr_nz_uv[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
        ]
        nxy_bini[direction] = curr_nxy_uv[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
        ]
        nz_bini_image[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            shift,
        ] = nz_bini[direction]
        nxy_bini_image[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == shift
            ],
            shift,
        ] = nxy_bini[direction]
    num_is_valid_pixel = v_where_is_valid.shape[0]
    num_is_valid_and_valid_neighbor = v_where_is_valid_and_valid_neighbor.shape[0]
    num_equations = num_is_valid_and_valid_neighbor

    assert H * W < 2**62 - 1
    valid_pixel_to_valid_idx = -np.ones([H, W], dtype=np.int64)
    valid_pixel_to_valid_idx[v_where_is_valid, u_where_is_valid] = np.arange(
        num_is_valid_pixel
    )
    valid_v_u_shift_to_valid_idx = -np.ones([H, W, num_shifts], dtype=np.int64)
    valid_v_u_shift_to_valid_idx[
        v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor,
    ] = num_is_valid_pixel + np.arange(num_is_valid_and_valid_neighbor)

    indices_z_a = valid_pixel_to_valid_idx[
        v_where_is_valid_and_valid_neighbor, u_where_is_valid_and_valid_neighbor
    ]
    indices_z_b = valid_pixel_to_valid_idx[
        v_where_is_valid_and_valid_neighbor
        + channel_idx_to_du_dv[shift_where_is_valid_and_valid_neighbor, 1],
        u_where_is_valid_and_valid_neighbor
        + channel_idx_to_du_dv[shift_where_is_valid_and_valid_neighbor, 0],
    ]
    assert not (np.any(indices_z_a == -1) or np.any(indices_z_b == -1)), (
        np.where(indices_z_a == -1),
        np.where(indices_z_b == -1),
    )

    assert len(indices_z_a) == len(indices_z_b) == num_equations

    assert len(indices_z_a) == len(indices_z_b)
    eqn_indices_z_a = np.arange(len(indices_z_a))
    eqn_indices_z_b = np.arange(len(indices_z_b))

    col_indices_bini = np.concatenate([indices_z_a, indices_z_b])
    row_indices_bini = np.concatenate(
        [
            # This accounts for the z_a and z_b terms.
            eqn_indices_z_a,
            eqn_indices_z_b,
        ]
    )
    all_values_bini = np.concatenate(
        [
            # This accounts for the z_a and z_b terms.
            -nz_bini_image[
                v_where_is_valid_and_valid_neighbor,
                u_where_is_valid_and_valid_neighbor,
                shift_where_is_valid_and_valid_neighbor,
            ],
            nz_bini_image[
                v_where_is_valid_and_valid_neighbor,
                u_where_is_valid_and_valid_neighbor,
                shift_where_is_valid_and_valid_neighbor,
            ],
        ]
    )
    A_bini = scipy.sparse.csr_matrix(
        (all_values_bini, (row_indices_bini, col_indices_bini)),
        shape=(num_equations, num_is_valid_pixel),
    )
    b_bini = -nxy_bini_image[
        v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor,
    ]

    return (
        A_bini,
        (indices_z_a, indices_z_b),
    )


def compute_residual_and_wuv_image(
    A_ours_log_with_gamma,
    log_z,
    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,
    sigmoid_k_value,
):
    if sigmoid_k_value != 2:
        print(
            f"\033[91mNOTE: Setting sigmoid `k` to {sigmoid_k_value} rather than 2."
            "\033[0m"
        )
    residual_ = A_ours_log_with_gamma @ log_z
    residual_image = np.zeros((H, W, num_shifts))
    residual_image[
        v_where_is_valid_and_valid_neighbor,
        u_where_is_valid_and_valid_neighbor,
        shift_where_is_valid_and_valid_neighbor,
    ] = residual_

    wuv_image = np.zeros((H, W, num_shifts))
    wuv_image[:] = np.nan
    for channel_idx in range(num_shifts // 2):
        opposite_channel_idx = channel_idx + num_shifts // 2
        if num_shifts == 8:
            assert opposite_channel_idx in [4, 5, 6, 7]
        elif num_shifts == 4:
            assert opposite_channel_idx in [2, 3]
        else:
            raise ValueError(f"Invalid number of channels: {num_shifts}.")
        curr_w_image = sigmoid(
            residual_image[..., channel_idx] ** 2
            - residual_image[..., opposite_channel_idx] ** 2,
            k=sigmoid_k_value,
        )
        wuv_image[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == channel_idx
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == channel_idx
            ],
            channel_idx,
        ] = (
            1
            - curr_w_image[
                v_where_is_valid_and_valid_neighbor[
                    shift_where_is_valid_and_valid_neighbor == channel_idx
                ],
                u_where_is_valid_and_valid_neighbor[
                    shift_where_is_valid_and_valid_neighbor == channel_idx
                ],
            ]
        )

        wuv_image[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == opposite_channel_idx
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == opposite_channel_idx
            ],
            opposite_channel_idx,
        ] = curr_w_image[
            v_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == opposite_channel_idx
            ],
            u_where_is_valid_and_valid_neighbor[
                shift_where_is_valid_and_valid_neighbor == opposite_channel_idx
            ],
        ]

    return residual_image, wuv_image