misc.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/>.

From https://github.com/xucao-42/bilateral_normal_integration/.
"""

import numpy as np


def move_left(mask):
    return np.pad(mask, ((0, 0), (0, 1)), "constant", constant_values=0)[
        :, 1:
    ]  # Shift the input mask array to the left by 1, filling the right edge with zeros.


def move_right(mask):
    return np.pad(mask, ((0, 0), (1, 0)), "constant", constant_values=0)[
        :, :-1
    ]  # Shift the input mask array to the right by 1, filling the left edge with zeros.


def move_top(mask):
    return np.pad(mask, ((0, 1), (0, 0)), "constant", constant_values=0)[
        1:, :
    ]  # Shift the input mask array up by 1, filling the bottom edge with zeros.


def move_bottom(mask):
    return np.pad(mask, ((1, 0), (0, 0)), "constant", constant_values=0)[
        :-1, :
    ]  # Shift the input mask array down by 1, filling the top edge with zeros.


def move_top_left(mask):
    return np.pad(mask, ((0, 1), (0, 1)), "constant", constant_values=0)[
        1:, 1:
    ]  # Shift the input mask array up and to the left by 1, filling the bottom and right edges with zeros.


def move_top_right(mask):
    return np.pad(mask, ((0, 1), (1, 0)), "constant", constant_values=0)[
        1:, :-1
    ]  # Shift the input mask array up and to the right by 1, filling the bottom and left edges with zeros.


def move_bottom_left(mask):
    return np.pad(mask, ((1, 0), (0, 1)), "constant", constant_values=0)[
        :-1, 1:
    ]  # Shift the input mask array down and to the left by 1, filling the top and right edges with zeros.


def move_bottom_right(mask):
    return np.pad(mask, ((1, 0), (1, 0)), "constant", constant_values=0)[
        :-1, :-1
    ]  # Shift the input mask array down and to the right by 1, filling the top and left edges with zeros.


def construct_facets_from(mask):
    # Initialize an array 'idx' of the same shape as 'mask' with integers
    # representing the indices of valid pixels in the mask.

    idx = np.zeros_like(mask, dtype=int)
    idx[mask] = np.arange(np.sum(mask))

    # Generate masks for neighboring pixels to define facets

    facet_move_top_mask = move_top(mask)
    facet_move_left_mask = move_left(mask)
    facet_move_top_left_mask = move_top_left(mask)

    # Identify the top-left pixel of each facet by performing a logical AND operation
    # on the masks of neighboring pixels and the input mask.

    facet_top_left_mask = np.logical_and.reduce(
        (facet_move_top_mask, facet_move_left_mask, facet_move_top_left_mask, mask)
    )

    # Create masks for the other three vertices of each facet by shifting the top-left 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 a numpy array of facets by stacking the indices of the four vertices
    # of each facet along the last dimension. Each row of the resulting array represents
    # a single facet with the format [4, idx_top_left, idx_bottom_left, idx_bottom_right, idx_top_right].

    return np.stack(
        (
            4 * np.ones(np.sum(facet_top_left_mask)),
            idx[facet_top_left_mask],
            idx[facet_bottom_left_mask],
            idx[facet_bottom_right_mask],
            idx[facet_top_right_mask],
        ),
        axis=-1,
    ).astype(int)


def sigmoid(x, k=1):
    return 1 / (1 + np.exp(-k * x))