Zernike transform prototype

Author: carterbox

src/tike/zernike.py

@@ -78,7 +78,7 @@ def _bino(a: int, b: int) -> int:
     return result
 
 
-def basis(size: int, degree: int) -> np.array:
+def basis(size: int, degree_min: int, degree_max: int) -> np.array:
     """Return all circular Zernike basis up to given radial degree.
 
     Parameters
@@ -102,16 +102,16 @@ def basis(size: int, degree: int) -> np.array:
     theta = np.arctan2(coords[0], coords[1])
 
     basis = []
-    for m, n in valid_indices(degree):
+    for m, n in valid_indices(degree_min, degree_max):
         basis.append(Z(m, n, radius, theta))
 
     basis = np.stack(basis, axis=0)
     return basis
 
 
-def valid_indices(degree: int) -> tuple:
+def valid_indices(degree_min: int, degree_max: int) -> tuple:
     """Enumerate all valid zernike indices (m,n) up to the given degree."""
-    for n in range(0, degree):
+    for n in range(degree_min, degree_max):
         for m in range(-n, n + 1):
             if (n - abs(m)) % 2 == 0:
                 yield m, n

tests/data/probe.png

@@ -2,6 +2,7 @@
 import os
 
 import tike.zernike
+import tike.linalg
 import tike.view
 import matplotlib.pyplot as plt
 import numpy as np
@@ -50,3 +51,47 @@ def test_radial_1(self):
 
     def test_radial_2(self):
         self._radial_template(2)
+
+    def test_transform(self):
+
+        fname = os.path.join(testdir, 'result', 'zernike')
+        os.makedirs(fname, exist_ok=True)
+
+        # import libimage
+
+        # f0 = libimage.load('cryptomeria', 256)
+        f0 = plt.imread(os.path.join(testdir, 'data', 'probe.png'))
+        size = f0.shape[-1]
+        plt.imsave(os.path.join(fname, 'basis-0.png'), f0)
+
+        f0 = f0.reshape(size * size, 1)
+
+        _basis = []
+
+        for d in range(0, 100):
+            _basis.append(
+                tike.zernike.basis(
+                    size,
+                    degree_min=d,
+                    degree_max=d + 1,
+                ))
+
+            basis = np.concatenate(_basis, axis=0)
+            print(f'degree {d} - {len(basis)}')
+            basis = np.moveaxis(basis, 0, -1)
+            basis = basis.reshape(size * size, -1)
+            # weight only pixels inside basis
+            w = (basis[..., 0] > 0).astype('float32')
+
+            # x = tike.linalg.lstsq(basis, f0, weights=w)
+            x, _, _, _ = np.linalg.lstsq(basis, f0, rcond=1e-4)
+
+            f1 = basis @ x
+            f1 = f1.reshape(size, size)
+            plt.imsave(os.path.join(fname, f'basis-{d:02d}.png'), f1)
+
+            plt.figure()
+            plt.title(f"basis weights for {d} degree polynomials")
+            plt.bar(list(range(x.size)), x.flatten())
+            plt.savefig(os.path.join(fname, f'basis-w-{d:02d}.png'))
+            plt.close()