src/saliency.py

# -*- coding: utf-8 -*-
from typing import Literal

import numpy as np
import cv2


class SaliencyMap:
    """
    SaliencyMap 顕著性マップを計算するクラスです。
    Example:
        from lib.saliency import SaliencyMap
        
        saliency = SaliencyMap("SpectralResidual")
        success, saliencyMap = saliency.computeSaliency(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 computeSaliency(self, image: np.ndarray):
        """
        入力画像から顕著性マップを作成します。

        Parameters:
            image: 入力画像

        Returns:
           bool:
               true: SaliencyMap computed, false:NG
           np.ndarray: 顕著性マップ               
        """
        # 画像の顕著性を計算します。
        return self.saliency.computeSaliency(image)

def convertColorMap(
    image: np.ndarray,
    saliencyMap: np.ndarray,
    colormap_name: Literal["jet", "hot", "turbo"] = "jet"):
    """
    顕著性マップをカラーマップに変換後に、入力画像に重ね合わせします。
    
    Parameters:
        image: 入力画像
        saliencyMap: 顕著性マップ
        colormap_name: カラーマップの種類

    Returns:
        np.ndarray: 重ね合わせた画像(RGBA形式)
    """
    colormaps = {"jet": cv2.COLORMAP_JET, "hot": cv2.COLORMAP_HOT, "turbo": cv2.COLORMAP_TURBO}
    if colormap_name not in colormaps:
        raise ValueError(colormap_name) 

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