PyWaveletsとOpenCVで画像のWavelet変換

やりたいこと

画像処理を最近学び始め、
画像に対してウェーブレット変換を行ってみたいなぁと思っていたところ、
下記の素晴らしい記事を見つけました。

mokemokechicken様：Pythonで画像をWavelet変換するサンプル

しかし、記事作者様のmokemokechicken様も述べている通り
「画像の縦横サイズが同じじゃないと上手くいかない」という課題がありました。

そこで今回はmokemokechicken様のサンプルコードを参考にさせて頂きながら、
OpenCVで画像を読み込むことで、「入力画像の縦横サイズがどのような場合でも」画像に対するウェーブレット変換ができるようにしました。
（カッコよく書いてますが、OpenCVで画像の読み込みしたくてコードをいじってたらうまいこと出来てしまったというオチです）

必要なもの

・PyWavelets
・numpy
・OpenCV

PyWaveletsに関する簡単な解説

URL : http://www.pybytes.com/pywavelets/ref/2d-dwt-and-idwt.html

pywt.dwt2という関数を用いることで、２次元配列（今回は画像）に対する
２次元ウェーブレット変換が簡単に行えますが、

複数階層の２次元ウェーブレット変換ができるように拡張された関数が、
pywt.wavedec2として用意されており、今回はこちらの関数を扱いました。

サンプルコード

今回はlenaに対してwavelet変換しています

import pywt
import numpy as np
import cv2

def image_normalization(src_img):
    """
    白飛び防止のための正規化処理
    cv2.imshowでwavelet変換された画像を表示するときに必要（大きい値を持つ画像の時だけ）
    """
    norm_img = (src_img - np.min(src_img)) / (np.max(src_img) - np.min(src_img))
    return norm_img

def merge_images(cA, cH_V_D):
    """numpy.array を４つ(左上、(右上、左下、右下))連結させる"""
    cH, cV, cD = cH_V_D
    cH = image_normalization(cH) # 外してもok
    cV = image_normalization(cV) # 外してもok
    cD = image_normalization(cD) # 外してもok
    cA = cA[0:cH.shape[0], 0:cV.shape[1]] # 元画像が2の累乗でない場合、端数ができることがあるので、サイズを合わせる。小さい方に合わせます。
    return np.vstack((np.hstack((cA,cH)), np.hstack((cV, cD)))) # 左上、右上、左下、右下、で画素をくっつける

def coeffs_visualization(cof):
    norm_cof0 = cof[0]
    norm_cof0 = image_normalization(norm_cof0) # 外してもok
    merge = norm_cof0
    for i in range(1, len(cof)):
        merge = merge_images(merge, cof[i])  # ４つの画像を合わせていく
    cv2.imshow('', merge)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def wavelet_transform_for_image(src_image, level, M_WAVELET="db1", mode="sym"):
    data = src_image.astype(np.float64)
    coeffs = pywt.wavedec2(data, M_WAVELET, level=level, mode=mode)
    return coeffs

if __name__ == "__main__":

    filename = 'lena.jpg'
    LEVEL = 3

    # 'haar', 'db', 'sym' etc...
    # URL: http://pywavelets.readthedocs.io/en/latest/ref/wavelets.html
    MOTHER_WAVELET = "db1"

    im = cv2.imread(filename)

    print('LEVEL :', LEVEL)
    print('MOTHER_WAVELET', MOTHER_WAVELET)
    print('original image size: ', im.shape)

    """
    各BGRチャネル毎に変換
    cv2.imreadはB,G,Rの順番で画像を吐き出すので注意
    """
    B = 0
    G = 1
    R = 2
    coeffs_B = wavelet_transform_for_image(im[:, :, B], LEVEL, M_WAVELET=MOTHER_WAVELET)
    coeffs_G = wavelet_transform_for_image(im[:, :, G], LEVEL, M_WAVELET=MOTHER_WAVELET)
    coeffs_R = wavelet_transform_for_image(im[:, :, R], LEVEL, M_WAVELET=MOTHER_WAVELET)

    coeffs_visualization(coeffs_B)
    # coeffs_visualization(coeffs_G)
    # coeffs_visualization(coeffs_R)

実行結果

元画像(lena.jpg)

wavelet係数を画像化したもの（Bチャンネル）

最後に

何か訂正箇所などありましたらご意見いただければと思います。