src/saliency.py

# -*- coding: utf-8 -*-
"""SaliencyMapを計算する"""
from typing import Any, Tuple, Literal

import numpy as np
import cv2


class SaliencyMap:
    """

    顕著性マップを計算するクラス。

    Example:

        from src.saliency import SaliencyMap



        saliency = SaliencyMap("SpectralResidual")

        success, saliencyMap = saliency.compute(image)

    """
    def __init__(

        self,

        algorithm: Literal["SpectralResidual", "FineGrained"] = "SpectralResidual",

    ):
        self.algorithm = algorithm
        # OpenCVのsaliencyを作成します。
        if algorithm == "SpectralResidual":
            self.saliency = cv2.saliency.StaticSaliencySpectralResidual_create()
        else:
            self.saliency = cv2.saliency.StaticSaliencyFineGrained_create()

    def compute(self, image: np.ndarray) -> Tuple[bool, Any]:
        """

        入力画像から顕著性マップを作成します。



        Parameters:

            image: 入力画像



        Returns:

           bool:

               true: SaliencyMap computed, false:NG

           np.ndarray: 顕著性マップ               

        """
        # 画像の顕著性を計算します。
        return self.saliency.computeSaliency(image)


def convert_colormap(

    image: np.ndarray,

    saliency_map: np.ndarray,

    colormap_name: Literal["jet", "hot", "turbo"] = "jet"

):
    """

    顕著性マップをカラーマップに変換後に、入力画像に重ね合わせします。



    Parameters:

        image: 入力画像

        saliency_map: 顕著性マップ

        colormap_name: カラーマップの種類



    Returns:

        np.ndarray: 重ね合わせた画像(RGBA形式)

    """
    maps = {"jet": cv2.COLORMAP_JET, "hot": cv2.COLORMAP_HOT, "turbo": cv2.COLORMAP_TURBO}
    if colormap_name not in maps:
        raise ValueError(colormap_name)

    # 顕著性マップをカラーマップに変換
    saliency_map = (saliency_map * 255).astype("uint8")
    saliency_map = cv2.applyColorMap(saliency_map, maps[colormap_name])
    #return saliencyMap
    # 入力画像とカラーマップを重ね合わせ
    overlay = cv2.addWeighted(image, 0.5, saliency_map, 0.5, 0)
    #return overlay
    return cv2.cvtColor(overlay, cv2.COLOR_BGR2RGBA)