codebase / stimlib

Commit a296b8336c3250b59c0d8847ff948276ef158b8c

Authored by David Mayerich
1 parent 9127f47e

added the muve alignment Python script

Showing 2 changed files with 161 additions and 324 deletions   Show diff stats
Inline Side-by-side
python/coupledwave.py deleted
View file @9127f47
1   -# -*- coding: utf-8 -*-
2   -"""
3   -Created on Thu Sep 6 22:14:36 2018
4   -
5   -@author: david
6   -"""
7   -
8   -import numpy as np
9   -import matplotlib.pyplot as plt
10   -
11   -#evaluate a plane wave for all points in R = [Rx, Ry, Rz]
12   -# E0 [Ex, Ey, Ez] is the electric field vector
13   -# k is the k vector
14   -# n is the refractive index of the material
15   -# R is a list of coordinates in x, y, and z
16   -def planewave2field(E0, k, n, R):
17   - knr = n * (k[0] * R[0] + k[1] * R[1] + k[2] + R[2])
18   - exp_2pknr = np.exp(1j * 2 * np.pi * knr)
19   - Ex = E0[0] * exp_2pknr
20   - Ey = E0[1] * exp_2pknr
21   - Ez = E0[2] * exp_2pknr
22   -
23   - return [Ex, Ey, Ez]
24   -
25   -def orthogonalize(E, k):
26   - s = np.cross(E, k)
27   - return np.cross(k, s)
28   -
29   -class layersample:
30   -
31   - #generate a layered sample based on a list of layer refractive indices
32   - # layer_ri is a list of complex refractive indices for each layer
33   - # z is the list of boundary positions along z
34   - def __init__(self, layer_ri, z):
35   - self.n = np.array(layer_ri).astype(np.complex128)
36   - self.z = np.array(z)
37   -
38   - #calculate the index of the field component associated with a layer
39   - # l is the layer index [0, L)
40   - # c is the component (x=0, y=1, z=2)
41   - # d is the direction (0 = transmission, 1 = reflection)
42   - def idx(self, l, c, d):
43   - i = l * 6 + d * 3 + c - 3
44   - print(i)
45   - return i
46   -
47   - #create a matrix for a single plane wave specified by k and E
48   - def solve(self, k, E):
49   -
50   - #calculate the wavenumber
51   - kn = np.linalg.norm(k)
52   -
53   - #s is the plane wave direction scaled by the refractive index
54   - s = k / kn * self.n[0]
55   -
56   - #allocate space for the matrix
57   - L = len(self.n)
58   - M = np.zeros((6*(L-1), 6*(L-1)), dtype=np.complex128)
59   -
60   - #allocate space for the RHS vector
61   - b = np.zeros(6*(L-1), dtype=np.complex128)
62   -
63   - #initialize a counter for the equation number
64   - ei = 0
65   -
66   - #calculate the sz component for each layer
67   - self.sz = np.zeros(L, dtype=np.complex128)
68   - for l in range(L):
69   - self.sz[l] = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
70   -
71   - #set constraints based on Gauss' law
72   - for l in range(0, L):
73   - #sz = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
74   -
75   - #set the upward components for each layer
76   - # note that layer L-1 does not have a downward component
77   - # David I, Equation 7
78   - if l != L-1:
79   - M[ei, self.idx(l, 0, 1)] = s[0]
80   - M[ei, self.idx(l, 1, 1)] = s[1]
81   - M[ei, self.idx(l, 2, 1)] = -self.sz[l]
82   - ei = ei+1
83   -
84   - #set the downward components for each layer
85   - # note that layer 0 does not have a downward component
86   - # Davis I, Equation 6
87   - if l != 0:
88   - M[ei, self.idx(l, 0, 0)] = s[0]
89   - M[ei, self.idx(l, 1, 0)] = s[1]
90   - M[ei, self.idx(l, 2, 0)] = self.sz[l]
91   - ei = ei+1
92   -
93   - #enforce a continuous field across boundaries
94   - for l in range(1, L):
95   - sz0 = self.sz[l-1]
96   - sz1 = self.sz[l]
97   - A = np.exp(1j * kn * sz0 * (self.z[l] - self.z[l-1]))
98   - if l < L - 1:
99   - B = np.exp(-1j * kn * sz1 * (self.z[l] - self.z[l+1]))
100   -
101   -
102   - #if this is the second layer, use the simplified equations that account for the incident field
103   - if l == 1:
104   - M[ei, self.idx(0, 0, 1)] = 1
105   - M[ei, self.idx(1, 0, 0)] = -1
106   - if L > 2:
107   - M[ei, self.idx(1, 0, 1)] = -B
108   -
109   - b[ei] = -A * E[0]
110   - ei = ei + 1
111   -
112   - M[ei, self.idx(0, 1, 1)] = 1
113   - M[ei, self.idx(1, 1, 0)] = -1
114   - if L > 2:
115   - M[ei, self.idx(1, 1, 1)] = -B
116   - b[ei] = -A * E[1]
117   - ei = ei + 1
118   -
119   - M[ei, self.idx(0, 2, 1)] = s[1]
120   - M[ei, self.idx(0, 1, 1)] = sz0
121   - M[ei, self.idx(1, 2, 0)] = -s[1]
122   - M[ei, self.idx(1, 1, 0)] = sz1
123   - if L > 2:
124   - M[ei, self.idx(1, 2, 1)] = -B*s[1]
125   - M[ei, self.idx(1, 1, 1)] = -B*sz1
126   - b[ei] = A * sz0 * E[1] - A * s[1]*E[2]
127   - ei = ei + 1
128   -
129   - M[ei, self.idx(0, 0, 1)] = -sz0
130   - M[ei, self.idx(0, 2, 1)] = -s[0]
131   - M[ei, self.idx(1, 0, 0)] = -sz1
132   - M[ei, self.idx(1, 2, 0)] = s[0]
133   - if L > 2:
134   - M[ei, self.idx(1, 0, 1)] = B*sz1
135   - M[ei, self.idx(1, 2, 1)] = B*s[0]
136   - b[ei] = A * s[0] * E[2] - A * sz0 * E[0]
137   - ei = ei + 1
138   -
139   - #if this is the last layer, use the simplified equations that exclude reflections from the last layer
140   - elif l == L-1:
141   - M[ei, self.idx(l-1, 0, 0)] = A
142   - M[ei, self.idx(l-1, 0, 1)] = 1
143   - M[ei, self.idx(l, 0, 0)] = -1
144   - ei = ei + 1
145   -
146   - M[ei, self.idx(l-1, 1, 0)] = A
147   - M[ei, self.idx(l-1, 1, 1)] = 1
148   - M[ei, self.idx(l, 1, 0)] = -1
149   - ei = ei + 1
150   -
151   - M[ei, self.idx(l-1, 2, 0)] = A*s[1]
152   - M[ei, self.idx(l-1, 1, 0)] = -A*sz0
153   - M[ei, self.idx(l-1, 2, 1)] = s[1]
154   - M[ei, self.idx(l-1, 1, 1)] = sz0
155   - M[ei, self.idx(l, 2, 0)] = -s[1]
156   - M[ei, self.idx(l, 1, 0)] = sz1
157   - ei = ei + 1
158   -
159   - M[ei, self.idx(l-1, 0, 0)] = A*sz0
160   - M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
161   - M[ei, self.idx(l-1, 0, 1)] = -sz0
162   - M[ei, self.idx(l-1, 2, 1)] = -s[0]
163   - M[ei, self.idx(l, 0, 0)] = -sz1
164   - M[ei, self.idx(l, 2, 0)] = s[0]
165   - ei = ei + 1
166   - #otherwise use the full set of boundary conditions
167   - else:
168   - M[ei, self.idx(l-1, 0, 0)] = A
169   - M[ei, self.idx(l-1, 0, 1)] = 1
170   - M[ei, self.idx(l, 0, 0)] = -1
171   - M[ei, self.idx(l, 0, 1)] = -B
172   - ei = ei + 1
173   -
174   - M[ei, self.idx(l-1, 1, 0)] = A
175   - M[ei, self.idx(l-1, 1, 1)] = 1
176   - M[ei, self.idx(l, 1, 0)] = -1
177   - M[ei, self.idx(l, 1, 1)] = -B
178   - ei = ei + 1
179   -
180   - M[ei, self.idx(l-1, 2, 0)] = A*s[1]
181   - M[ei, self.idx(l-1, 1, 0)] = -A*sz0
182   - M[ei, self.idx(l-1, 2, 1)] = s[1]
183   - M[ei, self.idx(l-1, 1, 1)] = sz0
184   - M[ei, self.idx(l, 2, 0)] = -s[1]
185   - M[ei, self.idx(l, 1, 0)] = sz1
186   - M[ei, self.idx(l, 2, 1)] = -B*s[1]
187   - M[ei, self.idx(l, 1, 1)] = -B*sz1
188   - ei = ei + 1
189   -
190   - M[ei, self.idx(l-1, 0, 0)] = A*sz0
191   - M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
192   - M[ei, self.idx(l-1, 0, 1)] = -sz0
193   - M[ei, self.idx(l-1, 2, 1)] = -s[0]
194   - M[ei, self.idx(l, 0, 0)] = -sz1
195   - M[ei, self.idx(l, 2, 0)] = s[0]
196   - M[ei, self.idx(l, 0, 1)] = B*sz1
197   - M[ei, self.idx(l, 2, 1)] = B*s[0]
198   - ei = ei + 1
199   -
200   - #store the matrix and RHS vector (for debugging)
201   - self.M = M
202   - self.b = b
203   -
204   - #evaluate the linear system
205   - P = np.linalg.solve(M, b)
206   -
207   - #save the results (also for debugging)
208   - self.P = P
209   -
210   - #store the coefficients for each layer
211   - self.Ptrans = []
212   - self.Prefl = []
213   - for l in range(L):
214   - if l == 0:
215   - self.Ptrans.append((E[0], E[1], E[2]))
216   - else:
217   - px = P[self.idx(l, 0, 0)]
218   - py = P[self.idx(l, 1, 0)]
219   - pz = P[self.idx(l, 2, 0)]
220   - self.Ptrans.append((px, py, pz))
221   -
222   - if l == L-1:
223   - self.Prefl.append((0, 0, 0))
224   - else:
225   - px = P[self.idx(l, 0, 1)]
226   - py = P[self.idx(l, 1, 1)]
227   - pz = P[self.idx(l, 2, 1)]
228   - self.Prefl.append((px, py, pz))
229   -
230   - #this converts the coefficients into a NumPy array for convenience later on
231   - self.Ptrans = np.array(self.Ptrans)
232   - self.Prefl = np.array(self.Prefl)
233   -
234   - #store values required for evaluation
235   - #store k
236   - self.k = kn
237   -
238   - #store sx and sy
239   - self.s = np.array([s[0], s[1]])
240   -
241   - self.solved = True
242   -
243   - #evaluate a solved homogeneous substrate
244   - def evaluate(self, X, Y, Z):
245   -
246   - if not self.solved:
247   - print("ERROR: the layered substrate hasn't been solved")
248   - return
249   -
250   - #this code is a bit cumbersome and could probably be optimized
251   - # Basically, it vectorizes everything by creating an image
252   - # such that the value at each pixel is the corresponding layer
253   - # that the pixel resides in. That index is then used to calculate
254   - # the field within the layer
255   -
256   - #allocate space for layer indices
257   - LI = np.zeros(Z.shape, dtype=np.int)
258   -
259   - #find the layer index for each sample point
260   - L = len(self.z)
261   - LI[Z < self.z[0]] = 0
262   - for l in range(L-1):
263   - idx = np.logical_and(Z > self.z[l], Z <= self.z[l+1])
264   - LI[idx] = l
265   - LI[Z > self.z[-1]] = L - 1
266   -
267   - #calculate the appropriate phase shift for the wave transmitted through the layer
268   - Ph_t = np.exp(1j * self.k * self.sz[LI] * (Z - self.z[LI]))
269   -
270   - #calculate the appropriate phase shift for the wave reflected off of the layer boundary
271   - LIp = LI + 1
272   - LIp[LIp >= L] = 0
273   - Ph_r = np.exp(-1j * self.k * self.sz[LI] * (Z - self.z[LIp]))
274   - Ph_r[LI >= L-1] = 0
275   -
276   - #calculate the phase shift based on the X and Y positions
277   - Ph_xy = np.exp(1j * self.k * (self.s[0] * X + self.s[1] * Y))
278   -
279   - #apply the phase shifts
280   - Et = self.Ptrans[LI] * Ph_t[:, :, None]
281   - Er = self.Prefl[LI] * Ph_r[:, :, None]
282   -
283   - #add everything together coherently
284   - E = (Et + Er) * Ph_xy[:, :, None]
285   -
286   - #return the electric field
287   - return E
288   -
289   -
290   -#This sample code produces a field similar to Figure 2 in Davis et al., 2010
291   -#set the material properties
292   -depths = [-50, -15, 15, 40]
293   -n = [1.0, 1.4+1j*0.05, 1.4, 1.0]
294   -
295   -#create a layered sample
296   -layers = layersample(n, depths)
297   -
298   -#set the input light parameters
299   -d = np.array([0.2, 0, 1])
300   -d = d / np.linalg.norm(d)
301   -l = 5.0
302   -k = 2 * np.pi / l * d
303   -E0 = [0, 1, 0]
304   -E0 = orthogonalize(E0, d)
305   -
306   -#solve for the substrate field
307   -layers.solve(k, E0)
308   -
309   -#set the simulation domain
310   -N = 512
311   -M = 1024
312   -D = [-50, 80]
313   -
314   -x = np.linspace(D[0], D[1], N)
315   -z = np.linspace(D[0], D[1], M)
316   -[X, Z] = np.meshgrid(x, z)
317   -Y = np.zeros(X.shape)
318   -E = layers.evaluate(X, Y, Z)
319   -
320   -Er = np.real(E)
321   -I = Er[..., 0] ** 2 + Er[..., 1] **2 + Er[..., 2] ** 2
322   -plt.figure()
323   -plt.set_cmap("afmhot")
324   -plt.imshow(Er[..., 1])
python/muve-align.py 0 → 100644
  Show/Hide comments   View file @a296b83
  1 +# -*- coding: utf-8 -*-
  2 +"""
  3 +Created on Fri May 4 17:15:07 2018
  4 +
  5 +@author: david
  6 +"""
  7 +
  8 +import imagestack
  9 +import numpy
  10 +import matplotlib.pyplot as plt
  11 +import cv2
  12 +import progressbar
  13 +import scipy.ndimage
  14 +import sys
  15 +
  16 +
  17 +def press(event):
  18 + global i
  19 + global ax
  20 + global fig
  21 + if event.key == 'z':
  22 + i = i - 1
  23 + if i < 0:
  24 + i = 0
  25 + if event.key == 'x':
  26 + i = i + 1
  27 + if i == S.shape[0]:
  28 + i = S.shape[0]
  29 + if event.key == "up":
  30 + for n in range(i, S.shape[0]):
  31 + warps[n][1, 2] = warps[n][1, 2] + 1
  32 + if event.key == "down":
  33 + for n in range(i, S.shape[0]):
  34 + warps[n][1, 2] = warps[n][1, 2] - 1
  35 + if event.key == "left":
  36 + for n in range(i, S.shape[0]):
  37 + warps[n][0, 2] = warps[n][0, 2] + 1
  38 + if event.key == "right":
  39 + for n in range(i, S.shape[0]):
  40 + warps[n][0, 2] = warps[n][0, 2] - 1
  41 + aligned = cv2.warpAffine(S[i, :, :, :], warps[i], (S.shape[2], S.shape[1]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP);
  42 + ax.cla()
  43 + ax.imshow(aligned.astype(numpy.uint8))
  44 + ax.set_title('Slice ' + str(i))
  45 + fig.canvas.draw()
  46 + sys.stdout.flush
  47 +
  48 +#apply the alignment adjustments to I based on the list of warp matrices in t
  49 +def apply_alignment(I, t):
  50 + A = numpy.zeros(I.shape, dtype=numpy.uint8)
  51 + for i in range(0, I.shape[0]):
  52 + A[i, :, :, :] = cv2.warpAffine(I[i, :, :, :], t[i], (I.shape[2], I.shape[1]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP);
  53 +
  54 + return A
  55 +
  56 +#get the transform that aligns image B to image A
  57 +def align(A, B, max_power=5):
  58 + # Define the motion model
  59 + warp_mode = cv2.MOTION_TRANSLATION
  60 +
  61 + # Specify the number of iterations.
  62 + number_of_iterations = 5000;
  63 +
  64 + # Specify the threshold of the increment
  65 + # in the correlation coefficient between two iterations
  66 + termination_eps = 1e-10;
  67 +
  68 + # Define termination criteria
  69 + criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations, termination_eps)
  70 +
  71 + #attempt to fit the image, increasing the blur as necessary
  72 + i = 0
  73 + while i <= max_power:
  74 + #blur the image used for fitting
  75 + im1_blur = scipy.ndimage.filters.gaussian_filter(A, (2 ** i, 2 ** i), mode='constant')
  76 + im2_blur = scipy.ndimage.filters.gaussian_filter(B, (2 ** i, 2 ** i), mode='constant')
  77 +
  78 + # Run the ECC algorithm. The results are stored in warp_matrix.
  79 + # Define 2x3 matrix and initialize the matrix to identity
  80 + warp_matrix = numpy.eye(2, 3, dtype=numpy.float32)
  81 + try:
  82 + (cc, warp_matrix) = cv2.findTransformECC (im1_blur,im2_blur,warp_matrix, warp_mode, criteria)
  83 + except:
  84 + #print("Error aligning at p = " + str(i))
  85 + i = i + 1
  86 + else:
  87 + #print("Successful alignment at p = " + str(i))
  88 + break
  89 + #enforce the fact that the x-axis is already aligned
  90 + #warp_matrix[0, 2] = 0
  91 + return warp_matrix
  92 +
  93 + #if i > 0:
  94 + # (cc, warp_matrix) = cv2.findTransformECC (A,B,warp_matrix, warp_mode, criteria)
  95 + #warp_matrix[0, 2] = 0
  96 + # return warp_matrix
  97 +
  98 +if len(sys.argv) < 3:
  99 + print("usage: muve-align input_mask output_directory")
  100 + print("ex: muve-align *.bmp ./aligned")
  101 + return
  102 +
  103 +fmask = sys.argv[1]
  104 +out_dir = sys.argv[2]
  105 +
  106 +if not os.path.isdir(out_dir):
  107 + os.mkdir(out_dir)
  108 +
  109 +#read the image stack for alignment
  110 +print("Loading image stack...")
  111 +S = imagestack.load(fmask, dtype=numpy.float32)
  112 +#convert to grayscale
  113 +G = imagestack.rgb2gray(S)
  114 +
  115 +# Define the motion model
  116 +warp_mode = cv2.MOTION_TRANSLATION
  117 +
  118 +# Define the motion model
  119 +warp_mode = cv2.MOTION_TRANSLATION
  120 +
  121 +# Specify the number of iterations.
  122 +number_of_iterations = 5000
  123 +
  124 +# Specify the threshold of the increment
  125 +# in the correlation coefficient between two iterations
  126 +termination_eps = 1e-10
  127 +
  128 +# Define termination criteria
  129 +criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations, termination_eps)
  130 +
  131 +print("\n\nCalculating alignment transformations...")
  132 +bar = progressbar.ProgressBar(max_value=G.shape[0])
  133 +#for each image in the stack, calculate a transformation matrix to align it to the previous image
  134 +warps = [numpy.eye(2, 3, dtype=numpy.float32)]
  135 +for i in range(1, G.shape[0]):
  136 + warps.append(align(G[i-1, :, :], G[i, :, :]))
  137 + bar.update(i+1)
  138 +
  139 +
  140 +print("\n\nApplying alignment transformation...")
  141 +# Define 2x3 matrix and initialize the matrix to identity
  142 +for i in range(1, G.shape[0]):
  143 + warps[i][0, 2] = warps[i][0, 2] + warps[i-1][0, 2]
  144 + warps[i][1, 2] = warps[i][1, 2] + warps[i-1][1, 2]
  145 +
  146 +
  147 +
  148 +i = 0
  149 +aligned = cv2.warpAffine(S[i, :, :, :], warps[i], (S.shape[2], S.shape[1]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
  150 +
  151 +fig, ax = plt.subplots()
  152 +
  153 +fig.canvas.mpl_connect('key_press_event', press)
  154 +
  155 +ax.imshow(aligned.astype(numpy.uint8))
  156 +ax.set_title('Slice ' + str(i))
  157 +ax.set_xlabel("Press 'z' and 'x' to change slices and arrow keys to align")
  158 +plt.show()
  159 +
  160 +I = apply_alignment(S, warps)
  161 +imagestack.save(I, out_dir + "/aligned_", ".bmp")
0 162 \ No newline at end of file
... ...