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 \ No newline at end of file 162 \ No newline at end of file