add mvdr implement and optimize cgmm-training process

Author: funcwj

.gitignore

@@ -1,6 +1,8 @@
 *.pyc
 *.log
 *.scp
+*.npy
+*.wav
 data/*.pyc
 __pycache__/
 6ch/

apply_mvdr.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+# coding=utf-8
+# wujian@17.10.27
+
+import argparse
+import numpy as np
+import beamformer
+import utils
+from utils import MultiChannelWrapper
+
+def main(args):
+    """
+        M: num_chanels, T: num_frames
+        apply_mvdr inputs:
+            steer_vector:     1 x M
+            sigma_noise[f]:      M x M
+            spectrum_onbin[f]:   T x M
+            return 1 x T
+    """
+    sigma_noisy = np.load(args.sigma_noisy)
+    sigma_noise = np.load(args.sigma_noise)
+    sigma_clean = sigma_noisy - sigma_noise
+
+    wrapper = MultiChannelWrapper(args.descriptor)
+    (time_steps, num_bins), spectrums = wrapper.spectrums()
+    specs_noisy = np.transpose(spectrums, (2, 1, 0)) 
+    specs_enhan = np.zeros([num_bins, time_steps]).astype(np.complex)
+    for f in range(num_bins):
+        steer_vector = beamformer.main_egvec(sigma_clean[f])
+        specs_enhan[f] = beamformer.apply_mvdr(steer_vector, sigma_noise[f], specs_noisy[f]) 
+    utils.reconstruct_wave(np.transpose(specs_enhan), args.save_dir, filter_coeff=args.filter_coeff)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Apply CGMM-MVDR beamformer on multiple channel")
+    parser.add_argument('descriptor', type=str,
+                        help="""descriptor of multiple channel location""")
+    parser.add_argument('sigma_noisy', type=str,
+                        help="""sigma of noisy(noise + clean) part estimated by CGMM""")
+    parser.add_argument('sigma_noise', type=str,
+                        help="""sigma of noise part estimated by CGMM""")
+    parser.add_argument('-s', '--save',
+                        dest='save_dir', type=str, default='default.wav',
+                        help="""path to save the enhanced wave""")
+    parser.add_argument('-c', '--filter_coeff',
+                        dest='filter_coeff', type=float, default='0.97',
+                        help="""filter coefficient to apply when reconstruct wave""")
+    args = parser.parse_args()
+    main(args)
+

beamformer.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# coding=utf-8
+# wujian@17.10.26
+
+import numpy as np
+
+def main_egvec(mat):
+    """
+        return the eigen vector as a estimate of steer vector, 
+        which has maximum eigen value
+    """
+    assert mat.ndim == 2, "Input must be 2-dim matrix/ndarray"
+    eigen_val, eigen_vec = np.linalg.eig(mat)
+    max_index = np.argsort(eigen_val)[-1]
+    return eigen_vec[max_index]
+
+def apply_mvdr(steer_vector, sigma_noise, spectrum_onbin):
+    """
+        inputs:
+            steer_vector:   M x 1 => d
+            sigma_noise:    M x M => \phi_v
+            spectrum_onbin: T x M => y
+        w = \phi_v^{-1} * d / (d^H * \phi_v^{-1} * d) => M x 1
+        s = w^H * y^T => 1 x T
+    """
+    # T x M => M x T
+    y = np.matrix(spectrum_onbin).T
+    # 1 x M => M x 1
+    d = np.matrix(steer_vector).T
+    phi_inv = np.matrix(sigma_noise).I
+    # M x 1
+    w = phi_inv * d / (d.H * phi_inv * d)
+    s = w.H * y
+    return s

cgmm.py

@@ -3,6 +3,7 @@
 # wujian@17.10.25
 
 import math
+import os
 import numpy as np
 
 LOG_PI = math.log(math.pi)
@@ -45,76 +46,90 @@ def __init__(self, num_bins, time_steps, num_channels):
         # lambda, phi, R for noisy/noise part
         self.lambda_ = np.random.rand(num_bins, time_steps).astype(np.complex)
         self.phi     = np.ones([num_bins, time_steps]).astype(np.complex)
-        # type matrix
-        self.R       = [np.matrix(np.eye(num_channels, num_channels).astype(np.complex)) \
-                            for i in range(num_bins)] 
-    
+
+    def init_sigma(self, sigma):
+        """
+            Inputs: sigma is a np.matrix list 
+            Keeps \sigma^{-1} and det(\sigma), \sigma equals \mean(y^H * y)
+        """
+        assert type(sigma) == list
+        self.sigma_inv = [mat.I for mat in sigma]
+        self.sigma_det = [np.linalg.det(mat) for mat in sigma]
+        
     def check_inputs(self, inputs):
         num_bins, time_steps, num_channels = inputs.shape
         assert num_bins == self.num_bins and time_steps == self.time_steps \
-            and num_channels == self.dim, 'inputs dim does not match CGMM config'
+            and num_channels == self.dim, 'Inputs dim does not match CGMM config'
 
     def log_likelihood(self, spectrums):
         self.check_inputs(spectrums)
         posteriors = 0.0
         for f in range(self.num_bins):
-            sigma_inv = self.R[f].I
-            sigma_det = np.linalg.det(self.R[f])
             for t in range(self.time_steps):
                 posteriors += self.lambda_[f, t] * gmm_posterior(spectrums[f, t], \
-                        self.phi[f, t], sigma_inv, sigma_det) 
+                        self.phi[f, t], self.sigma_inv[f], self.sigma_det[f]) 
         return posteriors
 
     def accu_stats(self, spectrums):
         self.check_inputs(spectrums)
         stats = np.zeros([self.num_bins, self.time_steps]).astype(np.complex)
         for f in range(self.num_bins):
-            sigma_inv = self.R[f].I
-            sigma_det = np.linalg.det(self.R[f])
             for t in range(self.time_steps):
                 stats[f, t] = gmm_posterior(spectrums[f, t], self.phi[f, t], \
-                        sigma_inv, sigma_det) 
+                        self.sigma_inv[f], self.sigma_det[f]) 
         return stats
 
     def update_lambda(self, spectrums, stats):
         print('update lambda...')
+        assert stats.shape == self.lambda_.shape
         for f in range(self.num_bins):
-            sigma_inv = self.R[f].I
-            sigma_det = np.linalg.det(self.R[f])
             for t in range(self.time_steps):
                 self.lambda_[f, t] = gmm_posterior(spectrums[f, t], self.phi[f, t], \
-                        sigma_inv, sigma_det) / stats[f, t]
+                        self.sigma_inv[f], self.sigma_det[f])
+        self.lambda_ = self.lambda_ / stats
 
-    def update_phi(self, spectrums):
+    def update_phi(self, covar):
         print('update phi...')
         for f in range(self.num_bins):
-            inv_R = self.R[f].I
             for t in range(self.time_steps):
-                y = np.matrix(spectrums[f, t])
-                self.phi[f, t] = np.trace(y.H * y * inv_R) / self.dim
+                self.phi[f, t] = np.trace(covar[f * self.time_steps + t] * self.sigma_inv[f])
+        self.phi = self.phi / self.dim
 
-    def update_R(self, spectrums):
+    def update_R(self, covar):
         print('update R...')
         for f in range(self.num_bins):
             sum_lambda = self.lambda_[f].sum()
-            self.R[f] = 0
+            R = np.matrix(np.zeros([self.dim, self.dim]).astype(np.complex))
             for t in range(self.time_steps):
-                y = np.matrix(spectrums[f, t])
-                self.R[f] += self.lambda_[f, t] * y.H * y / self.phi[f, t]
-            self.R[f] = self.R[f] / sum_lambda
+                R += self.lambda_[f, t] * covar[f * self.time_steps + t] / self.phi[f, t]
+            R = R / sum_lambda
+            self.sigma_inv[f] = R.I 
+            self.sigma_det[f] = np.linalg.det(R)
 
-    def update_parameters(self, spectrums, stats):
+    def update_parameters(self, spectrums, covar, stats):
         self.check_inputs(spectrums)
+        assert len(covar) == self.num_bins * self.time_steps and type(covar) == list
         self.update_lambda(spectrums, stats)
-        self.update_phi(spectrums)
-        self.update_R(spectrums)
+        self.update_phi(covar)
+        self.update_R(covar)
 
 class CGMMTrainer(object):
     def __init__(self, num_bins, time_steps, num_channels):
         self.noise_part = CGMM(num_bins, time_steps, num_channels)
         self.noisy_part = CGMM(num_bins, time_steps, num_channels)
         self.num_samples = num_bins * time_steps
-    
+
+    def init_sigma(self, spectrums):
+        # precompute the covariance matrix of each channel
+        print("initialize sigma...")
+        num_bins, time_steps, num_channels = spectrums.shape
+        self.covar = [y.H * y for y in [np.matrix(spectrums[f, t]) \
+                for f in range(num_bins) for t in range(time_steps)]]
+        self.noise_part.init_sigma([np.matrix(np.eye(num_channels, \
+                num_channels).astype(np.complex)) for f in range(num_bins)])
+        self.noisy_part.init_sigma([sum(self.covar[f * time_steps: \
+               (f + 1) * time_steps]) / time_steps for f in range(num_bins)])
+        
     def log_likelihood(self, spectrums):
         return (self.noise_part.log_likelihood(spectrums) + \
                 self.noisy_part.log_likelihood(spectrums)) / self.num_samples
@@ -125,14 +140,22 @@ def accu_stats(self, spectrums):
                 self.noise_part.accu_stats(spectrums)
 
     def update_parameters(self, spectrums, stats):
-        self.noise_part.update_parameters(spectrums, stats)
-        self.noisy_part.update_parameters(spectrums, stats)
-
+        self.noise_part.update_parameters(spectrums, self.covar, stats)
+        self.noisy_part.update_parameters(spectrums, self.covar, stats)
+
+    def save_param(self, dest):
+        sigma_ny = [mat.I for mat in self.noisy_part.sigma_inv]
+        sigma_ne = [mat.I for mat in self.noise_part.sigma_inv]
+        if not os.path.exists(dest):
+            os.mkdir(dest)
+        np.save(os.path.join(dest, 'sigma_noisy'), sigma_ny)
+        np.save(os.path.join(dest, 'sigma_noise'), sigma_ne)
+        
     def train(self, spectrums, iters=30):
+        self.init_sigma(spectrums)
         print('Likelihood: ({0.real:.5f}, {0.imag:.5f}i)'.format(self.log_likelihood(spectrums)))
         for it in range(1, iters + 1):
             stats = self.accu_stats(spectrums)
             self.update_parameters(spectrums, stats)
             print('epoch {0:2d}: Likelihood = ({1.real:.5f}, {1.imag:.5f}i)'.format(it, \
                     self.log_likelihood(spectrums)))
-

train_cgmm.py

@@ -3,27 +3,31 @@
 # wujian@17.10.26
 
 import argparse
+import time
 import numpy as np
 
 from utils import MultiChannelWrapper
 from cgmm import CGMMTrainer
 
 def train(args):
     wrapper = MultiChannelWrapper(args.descriptor)  
-    (time_steps, num_bins), spectrums = wrapper.spectrums(transpose=True)
-    num_bins, time_steps, num_channels = np.array(spectrums).shape
-    trainer = CGMMTrainer(num_bins, time_steps, num_channels)
-    trainer.train(spectrums, iters=args.iters)
+    (time_steps, num_bins), spectrums = wrapper.spectrums()
+    trainer = CGMMTrainer(num_bins, time_steps, len(spectrums))
+    start_time = time.time()
+    trainer.train(np.transpose(spectrums), iters=args.iters)
+    finish_time = time.time()
+    print('Total raining time: {:.3f}s'.format(finish_time - start_time))
+    trainer.save_param(args.save_dir)
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Training CGMM on multiple channel")
     parser.add_argument('descriptor', type=str,
-                        help="""descriptor of multiple channel location, format:
-                                /path/to/channel1
-                                ...
-                                /path/to/channeln""")
+                        help="""descriptor of multiple channel location""")
     parser.add_argument('-i', '--iters',
                         dest='iters', type=int, default='10',
                         help="""number of iterations to train""")
+    parser.add_argument('-s', '--save',
+                        dest='save_dir', type=str, default='',
+                        help="""directory to save sigma of CGMM""")
     args = parser.parse_args()
     train(args)

utils.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/isr/bin/env python
 # coding=utf-8
 # wujian@17.10.24
 
@@ -33,7 +33,7 @@ def pre_emphase(signal, filter_coeff=0.97):
     signal[0] -= filter_coeff * signal[0]
     return signal
 
-def compute_spectrum(wave_wrapper, window_type='hamming'):
+def compute_spectrum(wave_wrapper, transpose=False, window_type='hamming'):
     """
         Compute the DFT of each frames in the wrapper:
         1. default apply hamming-window on each frame
@@ -53,7 +53,7 @@ def compute_spectrum(wave_wrapper, window_type='hamming'):
     for index in range(num_frames):
         feature_in[: frame_size] = frames[index] * window 
         spectrum[index] = np.fft.rfft(feature_in)
-    return spectrum
+    return spectrum if not transpose else np.transpose(spectrum)
 
 def plot_spectrum(spectrum, frame_duration, title="samples.wav"):
     """
@@ -82,7 +82,7 @@ def write_wave(samples, frame_rate, dest):
     dest_wave = wave.open(dest, "wb")
     # 1 channel; int16 default
     dest_wave.setparams((1, 2, frame_rate, samples.size, 'NONE', 'not compressed'))
-    dest_wave.writeframes(samples.astype(np.int16))
+    dest_wave.writeframes(samples.astype(np.int16).tostring())
     print("1 channels; 2 bytes per sample; {num_samples} samples; " \
             "{frame_rate} samples per sec. OUT[{path}]".format(path=dest, \
             num_samples=samples.size, frame_rate=frame_rate))
@@ -166,9 +166,9 @@ def subframes(self, normalize=True):
         return shape_per_item, frames
 
     def spectrums(self, transpose=False):
-        spects = [compute_spectrum(wrapper) for wrapper in self.wrappers]
+        spects = [compute_spectrum(wrapper, transpose) for wrapper in self.wrappers]
         shape_per_item = check_status(spects)
-        return shape_per_item, (spects if not transpose else np.transpose(spects))
+        return shape_per_item, spects
 
     def __str__(self):
         return '\n'.join([str(wrapper) for wrapper in self.wrappers])