ImgWaveletPy/ImgWavelet.py at master · shouno/ImgWaveletPy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# -*- coding: utf-8 -*-
#
#

"""
Image Wavelet クラスは，2の冪サイズのイメージを与えて
そのウェーブレット表現を扱うためのクラス
"""

import numpy as np
import pywt

def rgb2gray(rgbimg):
    """Change RGB image to gray one"""
    return np.dot(rgbimg[...,:3], [0.299, 0.587, 0.114])


class ImageWavelet:
    def __init__( self, imgmat, mode='db1' ):
        self.img = imgmat
        self.mode = mode
        self.coeffs = pywt.wavedec2( self.img, self.mode )
        self.maxlvl = len( self.coeffs )
        self.nH = 2**(self.maxlvl-1)
        self.nV = 2**(self.maxlvl-1)

    def BaseMat( self ):
        coeffs = pywt.wavedec2( np.zeros( (self.nV, self.nH) ), self.mode )
        nl = self.nV*self.nH
        Phi = np.array( np.zeros( (nl, nl) ) )
        coeffs[0][0] = 1.
        img = pywt.waverec2( coeffs, self.mode )
        Phi[:,0] = img.reshape( (nl,) )
        coeffs[0][0] = 0.
        clm = 1

        for lv in range( 1, self.maxlvl ):
            cH, cV, cD = coeffs[lv]

            shp = cH.shape
            cnum = np.prod( shp )
            for i in range( cnum ):
                q = np.zeros( shp )
                q.reshape( (cnum,) )[i] = 1.
                coeffs[lv] = (q, cV, cD)
                img = pywt.waverec2( coeffs, self.mode ).reshape( (nl,) )
                Phi[:,clm] = img
                clm = clm + 1
                coeffs[lv] = ( cH, cV, cD )

            shp = cV.shape
            cnum = np.prod( shp )
            for i in range( cnum ):
                q = np.zeros( shp )
                q.reshape( (cnum,) )[i] = 1.
                coeffs[lv] = (cH, q, cD)
                img = pywt.waverec2( coeffs, self.mode ).reshape( (nl,) )
                Phi[:,clm] = img
                clm = clm + 1
                coeffs[lv] = ( cH, cV, cD )

            shp = cD.shape
            cnum = np.prod( shp )
            for i in range( cnum ):
                q = np.zeros( shp )
                q.reshape( (cnum,) )[i] = 1.
                coeffs[lv] = (cH, cV, q)
                img = pywt.waverec2( coeffs, self.mode ).reshape( (nl,) )
                Phi[:,clm] = img
                clm = clm + 1
                coeffs[lv] = ( cH, cV, cD )

        return Phi

    def Wv2coeff( self, img=None ):
        """
        Return array means coefficients for the Wavelet bases
        expressed as column vectors
        """
        if img is None:
            img = self.img
        else:
            if img.shape != self.img.shape:
                raise Exception( 'Image Sizes incompatible' )
        coeffs = pywt.wavedec2( img, self.mode )
        nl = self.nH * self.nV
        ret = np.zeros( nl )
        ret[0] = coeffs[0][0][0]
        idx = 1
        for lv in range( 1, self.maxlvl ):
            for cc in coeffs[lv]:  # (cH, cV, cD)
                crow, cclm = cc.shape
                ret[idx:idx+crow*cclm] = cc.reshape( (crow*cclm,) )
                idx = idx + crow*cclm
        return ret

    def Coeff2Wv( self, coeff ):
        """
        Interprete coeffs as corresponding wavelet structure
        and return image for the coeffs
        """
        nl = self.nH * self.nV
        assert len( coeff ) == nl
        self.coeffs[0][0][0] = coeff[0]
        idx = 1
        for lv in range( 1, self.maxlvl ):
            cH, cV, cD = self.coeffs[lv]

            crow, cclm = cH.shape
            clen = crow*cclm
            cHnew = coeff[idx:idx+clen].reshape( (crow, cclm) )
            idx = idx + clen

            crow, cclm = cV.shape
            clen = crow*cclm
            cVnew = coeff[idx:idx+clen].reshape( (crow, cclm) )
            idx = idx + clen

            crow, cclm = cD.shape
            clen = crow*cclm
            cDnew = coeff[idx:idx+clen].reshape( (crow, cclm) )
            idx = idx + clen

            self.coeffs[lv] = (cHnew, cVnew, cDnew)

        self.img = pywt.waverec2( self.coeffs, self.mode )
        return self.img