image-match/python/py/Lib/site-packages/kornia/geometry/subpix/nms.py

# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import annotations

import torch
import torch.nn.functional as F

from kornia.core import Module, Tensor, eye, pad, zeros


def _get_nms_kernel2d(kx: int, ky: int) -> Tensor:
    """Return neigh2channels conv kernel."""
    numel: int = ky * kx
    center: int = numel // 2
    weight = eye(numel)
    weight[center, center] = 0
    return weight.view(numel, 1, ky, kx)


def _get_nms_kernel3d(kd: int, ky: int, kx: int) -> Tensor:
    """Return neigh2channels conv kernel."""
    numel: int = kd * ky * kx
    center: int = numel // 2
    weight = eye(numel)
    weight[center, center] = 0
    return weight.view(numel, 1, kd, ky, kx)


class NonMaximaSuppression2d(Module):
    r"""Apply non maxima suppression to filter.

    Flag `minima_are_also_good` is useful, when you want to detect both maxima and minima, e.g. for DoG
    """

    kernel: Tensor

    def __init__(self, kernel_size: tuple[int, int]) -> None:
        super().__init__()
        self.kernel_size: tuple[int, int] = kernel_size
        self.padding: tuple[int, int, int, int] = self._compute_zero_padding2d(kernel_size)
        self.register_buffer("kernel", _get_nms_kernel2d(*kernel_size))

    @staticmethod
    def _compute_zero_padding2d(kernel_size: tuple[int, int]) -> tuple[int, int, int, int]:
        # TODO: This method is duplicated with some utility function on kornia.filters
        if not isinstance(kernel_size, tuple):
            raise AssertionError(type(kernel_size))
        if len(kernel_size) != 2:
            raise AssertionError(kernel_size)

        def pad(x: int) -> int:
            return (x - 1) // 2  # zero padding function

        ky, kx = kernel_size  # we assume a cubic kernel
        return pad(ky), pad(ky), pad(kx), pad(kx)

    def forward(self, x: Tensor, mask_only: bool = False) -> Tensor:
        if len(x.shape) != 4:
            raise AssertionError(x.shape)
        B, CH, H, W = x.size()
        # find local maximum values
        x_padded = pad(x, list(self.padding)[::-1], mode="replicate")
        B, CH, HP, WP = x_padded.size()

        neighborhood = F.conv2d(x_padded.view(B * CH, 1, HP, WP), self.kernel.to(x.device, x.dtype), stride=1).view(
            B, CH, -1, H, W
        )
        max_non_center = neighborhood.max(dim=2)[0]
        mask = x > max_non_center
        if mask_only:
            return mask
        return x * (mask.to(x.dtype))


class NonMaximaSuppression3d(Module):
    r"""Apply non maxima suppression to filter."""

    def __init__(self, kernel_size: tuple[int, int, int]) -> None:
        super().__init__()
        self.kernel_size: tuple[int, int, int] = kernel_size
        self.padding: tuple[int, int, int, int, int, int] = self._compute_zero_padding3d(kernel_size)
        self.kernel = _get_nms_kernel3d(*kernel_size)

    @staticmethod
    def _compute_zero_padding3d(kernel_size: tuple[int, int, int]) -> tuple[int, int, int, int, int, int]:
        # TODO: This method is duplicated with some utility function on kornia.filters
        if not isinstance(kernel_size, tuple):
            raise AssertionError(type(kernel_size))
        if len(kernel_size) != 3:
            raise AssertionError(kernel_size)

        def pad(x: int) -> int:
            return (x - 1) // 2  # zero padding function

        kd, ky, kx = kernel_size  # we assume a cubic kernel
        return pad(kd), pad(kd), pad(ky), pad(ky), pad(kx), pad(kx)

    def forward(self, x: Tensor, mask_only: bool = False) -> Tensor:
        if len(x.shape) != 5:
            raise AssertionError(x.shape)
        # find local maximum values
        B, CH, D, H, W = x.size()
        if self.kernel_size == (3, 3, 3):
            mask = zeros(B, CH, D, H, W, device=x.device, dtype=torch.bool)
            center = slice(1, -1)
            left = slice(0, -2)
            right = slice(2, None)
            center_tensor = x[..., center, center, center]
            mask[..., 1:-1, 1:-1, 1:-1] = (
                (center_tensor > x[..., center, center, left])
                & (center_tensor > x[..., center, center, right])
                & (center_tensor > x[..., center, left, center])
                & (center_tensor > x[..., center, left, left])
                & (center_tensor > x[..., center, left, right])
                & (center_tensor > x[..., center, right, center])
                & (center_tensor > x[..., center, right, left])
                & (center_tensor > x[..., center, right, right])
                & (center_tensor > x[..., left, center, center])
                & (center_tensor > x[..., left, center, left])
                & (center_tensor > x[..., left, center, right])
                & (center_tensor > x[..., left, left, center])
                & (center_tensor > x[..., left, left, left])
                & (center_tensor > x[..., left, left, right])
                & (center_tensor > x[..., left, right, center])
                & (center_tensor > x[..., left, right, left])
                & (center_tensor > x[..., left, right, right])
                & (center_tensor > x[..., right, center, center])
                & (center_tensor > x[..., right, center, left])
                & (center_tensor > x[..., right, center, right])
                & (center_tensor > x[..., right, left, center])
                & (center_tensor > x[..., right, left, left])
                & (center_tensor > x[..., right, left, right])
                & (center_tensor > x[..., right, right, center])
                & (center_tensor > x[..., right, right, left])
                & (center_tensor > x[..., right, right, right])
            )
        else:
            max_non_center = (
                F.conv3d(
                    pad(x, list(self.padding)[::-1], mode="replicate"),
                    self.kernel.repeat(CH, 1, 1, 1, 1).to(x.device, x.dtype),
                    stride=1,
                    groups=CH,
                )
                .view(B, CH, -1, D, H, W)
                .max(dim=2, keepdim=False)[0]
            )
            mask = x > max_non_center
        if mask_only:
            return mask
        return x * (mask.to(x.dtype))


# functional api


def nms2d(input: Tensor, kernel_size: tuple[int, int], mask_only: bool = False) -> Tensor:
    r"""Apply non maxima suppression to filter.

    See :class:`~kornia.geometry.subpix.NonMaximaSuppression2d` for details.
    """
    return NonMaximaSuppression2d(kernel_size)(input, mask_only)


def nms3d(input: Tensor, kernel_size: tuple[int, int, int], mask_only: bool = False) -> Tensor:
    r"""Apply non maxima suppression to filter.

    See
    :class: `~kornia.feature.NonMaximaSuppression3d` for details.
    """
    return NonMaximaSuppression3d(kernel_size)(input, mask_only)