122 lines
3.7 KiB
Python
122 lines
3.7 KiB
Python
#!/usr/bin/python
|
|
# Coding: UTF-8
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
from PyEMD.BEMD import BEMD
|
|
except (ImportError, ModuleNotFoundError):
|
|
# Not supported until supported.
|
|
pass
|
|
|
|
|
|
@unittest.skip("Not supported until supported")
|
|
class BEMDTest(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.bemd = BEMD()
|
|
|
|
@staticmethod
|
|
def _generate_image(r=64, c=64):
|
|
return np.random.random((r, c))
|
|
|
|
@staticmethod
|
|
def _generate_linear_image(r=16, c=16):
|
|
rows = np.arange(r)
|
|
return np.repeat(rows, c).reshape(r, c)
|
|
|
|
@staticmethod
|
|
def _generate_Gauss(x, y, pos, std, amp=1):
|
|
x_s = x - pos[0]
|
|
y_s = y - pos[1]
|
|
x2 = x_s * x_s
|
|
y2 = y_s * y_s
|
|
exp = np.exp(-(x2 + y2) / (2 * std * std))
|
|
# exp[exp<1e-6] = 0
|
|
scale = amp / np.linalg.norm(exp)
|
|
return scale * exp
|
|
|
|
@classmethod
|
|
def _sin(cls, x_n=128, y_n=128, x_f=[1], y_f=[0], dx=0, dy=0):
|
|
x = np.linspace(0, 1, x_n) - dx
|
|
y = np.linspace(0, 1, y_n) - dy
|
|
xv, yv = np.meshgrid(x, y)
|
|
img = np.zeros(xv.shape)
|
|
for f in x_f:
|
|
img += np.sin(f * 2 * np.pi * xv)
|
|
for f in y_f:
|
|
img += np.cos(f * 2 * np.pi * yv)
|
|
return 255 * (img - img.min()) / (img.max() - img.min())
|
|
|
|
def test_extract_maxima(self):
|
|
image = self._sin(x_n=32, y_n=32, y_f=[1], dy=1)
|
|
max_peak_x, max_peak_y = BEMD.extract_maxima_positions(image)
|
|
self.assertEqual(max_peak_x.size, 6) # Clustering
|
|
self.assertEqual(max_peak_y.size, 6) # Clustering
|
|
|
|
def test_extract_minima(self):
|
|
image = self._sin(x_n=64, y_n=64, y_f=[2])
|
|
min_peak_x, min_peak_y = BEMD.extract_minima_positions(image)
|
|
self.assertEqual(min_peak_x.size, 16) # Clustering
|
|
self.assertEqual(min_peak_y.size, 16) # Clustering
|
|
|
|
def test_find_extrema(self):
|
|
image = self._sin()
|
|
min_peaks, max_peaks = BEMD.find_extrema_positions(image)
|
|
self.assertTrue(isinstance(min_peaks, tuple))
|
|
self.assertTrue(isinstance(min_peaks[0], np.ndarray))
|
|
self.assertTrue(isinstance(max_peaks[1], np.ndarray))
|
|
|
|
def test_default_call_BEMD(self):
|
|
x = np.arange(50)
|
|
y = np.arange(50)
|
|
xv, yv = np.meshgrid(x, y)
|
|
img = self._generate_Gauss(xv, yv, (10, 20), 5)
|
|
img += self._sin(x_n=x.size, y_n=y.size)
|
|
|
|
max_imf = 2
|
|
self.bemd(img, max_imf)
|
|
|
|
def test_endCondition_perfectReconstruction(self):
|
|
c1 = self._generate_image()
|
|
c2 = self._generate_image()
|
|
IMFs = np.stack((c1, c2))
|
|
org_img = np.sum(IMFs, axis=0)
|
|
self.assertTrue(self.bemd.end_condition(org_img, IMFs))
|
|
|
|
def test_bemd_simpleIMF(self):
|
|
image = self._sin(x_f=[3, 5], y_f=[2])
|
|
IMFs = self.bemd(image)
|
|
|
|
# One of the reasons this algorithm isn't the preferred one
|
|
self.assertTrue(IMFs.shape[0] == 7, "Depending on spline, there should be an IMF and possibly trend")
|
|
|
|
def test_bemd_limitImfNo(self):
|
|
|
|
# Create image
|
|
rows, cols = 64, 64
|
|
linear_background = 0.2 * self._generate_linear_image(rows, cols)
|
|
|
|
# Sinusoidal IMF
|
|
X = np.arange(cols)[None, :].T
|
|
Y = np.arange(rows)
|
|
x_comp_1d = np.sin(X * 0.3) + np.cos(X * 2.9) ** 2
|
|
y_comp_1d = np.sin(Y * 0.2)
|
|
comp_2d = 10 * x_comp_1d * y_comp_1d
|
|
comp_2d = comp_2d
|
|
|
|
image = linear_background + comp_2d
|
|
|
|
# Limit number of IMFs
|
|
max_imf = 2
|
|
|
|
# decompose image
|
|
IMFs = self.bemd(image, max_imf=max_imf)
|
|
|
|
# It should have no more than 2 (max_imf) + residue
|
|
self.assertEqual(IMFs.shape[0], 1 + max_imf)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|