grid.py 21.2 KB
Newer Older
1
import numpy
Edward Andò's avatar
Edward Andò committed
2
import spam.DIC.correlate
3
import spam.DIC.transformationOperator
4
import Queue
Edward Andò's avatar
Edward Andò committed
5
##import spam.DIC
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

def makeGrid( imageSize, nodeSpacing ):
    """
        Define a grid of correlation points.

        Parameters
        ----------
        imageSize : 3-item list
            Size of volume to spread the grid inside

        nodeSpacing : 3-item list or int
            Spacing between nodes

        Returns
        -------
21
        nodePositions : nPointsx3 numpy.array
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
            Array containing Z, Y, X positions of each point in the grid
    """
    if len( imageSize ) != 3:
        print( "\tgrid.makeGrid(): imageSize doesn't have three dimensions, exiting" )
        return

    if type( nodeSpacing ) == int or type( nodeSpacing ) == float:
        nodeSpacing = [ nodeSpacing ]*3
    elif len( nodeSpacing) != 3:
        print( "\tgrid.makeGrid(): nodeSpacing is not an int or float and doesn't have three dimensions, exiting" )
        return

    # Note: in this cheap node spacing, the first node is always at a distance of --nodeSpacing-- from the origin
    # The following could just be done once in principle...
    nodesMgrid = numpy.mgrid[ nodeSpacing[0]:imageSize[0]:nodeSpacing[0],
                              nodeSpacing[1]:imageSize[1]:nodeSpacing[1],
                              nodeSpacing[2]:imageSize[2]:nodeSpacing[2] ]
    nodesDim = (nodesMgrid.shape[1],nodesMgrid.shape[2],nodesMgrid.shape[3])

    numberOfNodes = int( nodesMgrid.shape[1] * nodesMgrid.shape[2] * nodesMgrid.shape[3] )

    nodePositions = numpy.zeros( ( numberOfNodes, 3 ) )

    nodePositions[:,0] = nodesMgrid[0].ravel()
    nodePositions[:,1] = nodesMgrid[1].ravel()
    nodePositions[:,2] = nodesMgrid[2].ravel()

Edward Andò's avatar
Edward Andò committed
49
    return nodePositions, nodesDim
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83




def pixelSearch( im1, im2, nodePositions, halfWindowSize, searchRange,
                 Ffield = None,
                 minSubVolume=None, im1mask=None,
                 greyLowThreshold = -numpy.inf):
    """
        This function handles grid-based local correlation, offering an initial rough dispalcement-only guess.
        At the moment matching of windows is done with a Normalised-Correlation-Coefficient approach.

        Parameters
        ----------
        im1 : 3D numpy array
            A 3D image of greylevels defining a reference configuration for the pixel search

        im2 : 3D numpy array
            A deformed 3D image of greylevels

        nodePositions : nPoints*3 numpy array
            Array containing Z, Y, X positions of each point in the grid, as returned by ``makeGrid`` for example

        halfWindowSize : 3-item list or int
            Size of subvolumes to perform the image correlation on, as a data range taken either side of the voxel on which the node is placed.
            The subvolume will be 2*halfWindowSize + 1 pixels on each side.
            A general recommendation is to make this half the node spacing

        searchRange : dictionary
            Search range as a dictionary containing 3 keys: 'zRange', 'yRange', and 'xRange',
            Each of which contains a list with two items

        Ffield : nPoints*4*4 numpy array, optional
            Optional field of ``F`` transformation operators defined for each node.
84
85
            Currently, only the translational components of F will be taken into account.
            Default = No displacement
86
87
88
89
90
91
92

        minSubVolume : int, optional
            Minimum number of pixels in a subvolume for it to be correlated (only considered in the case of im1mask).
            Expressed as a % of subvolume volume.
            Default = 50% of subvolume

        im1mask : 3D boolean numpy array, optional
93
94
            A mask for im1 which is true in the zones to correlate.
            Default = None
95
96
97

        greyLowThreshold : float, optional
            Threshold for the mean greylevel in each im1 subvolume.
98
99
            If the mean is below this value, the grid point is not correlated.
            Default = -inf
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

        Returns
        -------
        Dictionary containing:

            Keys
            ----
            Ffield : nNodes*4*4 numpy array of floats
                For each node, the measured transformatio operator (displacement only)

            pixelSearchCC : nNodes numpy array of floats
                For each node, the NCC score obtained
    """
    numberOfNodes = nodePositions.shape[0]

    # Check input sanity
    if type( halfWindowSize ) == int or type( halfWindowSize ) == float:
        halfWindowSize = [ halfWindowSize ]*3

    # Check minSubVolume
    if minSubVolume is None:
        minSubVolume = int( ((1+min(halfWindowSize)*2)**3) * 0.5 )
    else:
        minSubVolume = int( ((1+min(halfWindowSize)*2)**3) * minSubVolume )

    # Check F field
    if Ffield is None:
        Ffield = numpy.zeros( ( numberOfNodes, 4, 4 ) )
        for nodeNumber in xrange( numberOfNodes ): Ffield[nodeNumber] = numpy.eye( 4 )
129
130
131
132
133
134
135
136
137
    else:
        # Add initial displacement guess to search range
        for nodeNumber in xrange( numberOfNodes ): 
            searchRange['zRange'][0] += Ffield[nodeNumber,0,3]
            searchRange['zRange'][1] += Ffield[nodeNumber,0,3]
            searchRange['yRange'][0] += Ffield[nodeNumber,1,3]
            searchRange['yRange'][1] += Ffield[nodeNumber,1,3]
            searchRange['xRange'][0] += Ffield[nodeNumber,2,3]
            searchRange['xRange'][1] += Ffield[nodeNumber,2,3]
138
139
140
141
142
143
144
145
146
147

    # Create pixelSearchCC vector
    pixelSearchCC = numpy.zeros( ( numberOfNodes ) )

    print( "\tStarting Pixel search" )

    # Start loop over nodes
    for nodeNumber in xrange( numberOfNodes ):
        print( "\r\t\tCorrelating node {:04d} of {:04d}".format( nodeNumber+1, numberOfNodes ) ), 

148
149
150
151
        # Prepare slice for im1 -- this is always in the same place
        subVolSlice1 = [slice(int(nodePositions[nodeNumber,0]-halfWindowSize[0]), int(nodePositions[nodeNumber,0]+halfWindowSize[0]+1) ),
                        slice(int(nodePositions[nodeNumber,1]-halfWindowSize[1]), int(nodePositions[nodeNumber,1]+halfWindowSize[1]+1) ),
                        slice(int(nodePositions[nodeNumber,2]-halfWindowSize[2]), int(nodePositions[nodeNumber,2]+halfWindowSize[2]+1) ) ]
152
153
154
155
156

        # Extract it...
        imagette1 = im1[subVolSlice1].copy()

        # Check Mask volume condition
157
        if im1mask is not None: maskVolCondition = im1mask[subVolSlice1].sum() > minSubVolume
158
159
160
161
        else:                   maskVolCondition = True

        # Check it satisfies the volume and greyscale criteria
        if numpy.nanmean( imagette1 ) > greyLowThreshold and maskVolCondition:
Edward Andò's avatar
Edward Andò committed
162
            returns = spam.DIC.correlate.pixelSearch(   imagette1, im2,
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
                                                                searchRange  = searchRange,
                                                                searchCentre = nodePositions[ nodeNumber, 0:3 ] )
            Ffield[ nodeNumber, 0:3, 3 ] += returns['transformation']['t']
            pixelSearchCC[ nodeNumber ]   = returns['cc']

        # Either not enough data or below the threshold
        else:
            Ffield[ nodeNumber, 0:3, 3 ] = numpy.NaN
            pixelSearchCC[ nodeNumber ]  = numpy.NaN

    return { 'Ffield': Ffield,
             'pixelSearchCC': pixelSearchCC}





180
181
182
183
184
185
186
187
def subPixel(   im1, im2,
                nodePositions, halfWindowSize,
                Ffield = None,
                margin = None,
                maxIterations = None,
                minFchange = None,
                interpolationOrder = None,
                minSubVolume=None, im1mask=None,
188
189
                greyLowThreshold = -numpy.inf,
                mpi = False):
190
    """
191
        This function handles grid-based local correlation, performing a "lucasKanade" subpixel refinement.
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
        Here we minimise a residual which is the difference between im1 and im2.

        Parameters
        ----------
        im1 : 3D numpy array
            A 3D image of greylevels defining a reference configuration for the pixel search

        im2 : 3D numpy array
            A deformed 3D image of greylevels

        nodePositions : nPoints*3 numpy array
            Array containing Z, Y, X positions of each point in the grid, as returned by ``makeGrid`` for example

        halfWindowSize : 3-item list or int
            Size of subvolumes to perform the image correlation on, as a data range taken either side of the voxel on which the node is placed.
            The subvolume will be 2*halfWindowSize + 1 pixels on each side.
            A general recommendation is to make this half the node spacing

210
        Ffield : nPoints*4*4 numpy array, optional (default = No displacement)
211
212
            Optional field of ``F`` transformation operators defined for each node.

213
        margin : int, optional (default = None (use ``lucasKanade``'s default))
214
215
            Margin to extract for subpixel interpolation.

216
        maxIterations : int, optional (default = None (use ``lucasKanade``'s default))
217
218
            Number of iterations for subpixel refinement.

219
        minFchange : float, optional (default = None (use ``lucasKanade``'s default))
220
221
            Stop iterating when norm of F gets below this value.

222
        interpolationOrder : int, optional (default = None (use ``lucasKanade``'s default))
223
224
            Greyscale interpolation order.

225
        minSubVolume : int, optional (default = 50% of subvolume)
226
227
228
            Minimum number of pixels in a subvolume for it to be correlated (only considered in the case of im1mask).
            Expressed as a % of subvolume volume.

229
        im1mask : 3D boolean numpy array, optional (default = None)
230
231
            A mask for im1 which is true in the zones to correlate.

232
        greyLowThreshold : float, optional (default = -inf)
233
234
            Threshold for the mean greylevel in each im1 subvolume.
            If the mean is below this value, the grid point is not correlated.
235
236
237
            
        mpi : bool, optional (default = False)
            Are we being called by an MPI run?
238
    """
239
    
240
241
242
    if interpolationOrder == 1: interpolator = 'C'
    else:                       interpolator = 'python'
    
243
244
245
    # Add 1 to margin in order to correlate on the exact half-windo size asked by user
    margin += 1
    
246
    numberOfNodes = nodePositions.shape[0]
247
248
    
    def getImagettes( nodeNumber, im1, im2, Ffield, nodePositions, margin, halfWindowSize, greyLowThreshold, im1mask):
249
250
        # Initialise defaults
        imagette1mask = None
251
252
        Finit = None
        
Edward Andò's avatar
Edward Andò committed
253
254
255
        # Make sure all components of F are real numbers
        if numpy.isfinite( Ffield[ nodeNumber ] ).sum() == 16:
            # Apply this point's F (at the point's origin, i.e., the top HWS away from point)
256
            nodeDisplacement = numpy.round(Ffield[nodeNumber][0:3,-1])            
Edward Andò's avatar
Edward Andò committed
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292

            # Prepare im1 imagette -- and move it
            subVolSlice1 = [slice( int(nodePositions[nodeNumber,0] - halfWindowSize[0] - nodeDisplacement[0] - margin ), int( nodePositions[nodeNumber,0] + halfWindowSize[0] - nodeDisplacement[0] + margin + 1) ),
                            slice( int(nodePositions[nodeNumber,1] - halfWindowSize[1] - nodeDisplacement[1] - margin ), int( nodePositions[nodeNumber,1] + halfWindowSize[1] - nodeDisplacement[1] + margin + 1) ),
                            slice( int(nodePositions[nodeNumber,2] - halfWindowSize[2] - nodeDisplacement[2] - margin ), int( nodePositions[nodeNumber,2] + halfWindowSize[2] - nodeDisplacement[2] + margin + 1) ) ]

            # Extract it
            imagette1 = im1[subVolSlice1].copy()

            # If there is a mask defined

            # Check Mask volume condition
            if im1mask is not None:
                imagette1mask = im1mask[subVolSlice1]
                maskVolCondition = imagette1mask.sum() > minSubVolume
            else:
                imagette1mask = None
                maskVolCondition = True
                
            # Check if out extracted imagette 1 is above grey threshold, if there is enough under the mask if defined, and that dispalcements are not NaN
            if numpy.nanmean( imagette1 ) > greyLowThreshold and maskVolCondition:

                # prepare slice for imagette 2 -- this is fixed
                subVolSlice2 = [slice( int(nodePositions[nodeNumber,0] - halfWindowSize[0]), int(nodePositions[nodeNumber,0] + halfWindowSize[0]+1) ),
                                slice( int(nodePositions[nodeNumber,1] - halfWindowSize[1]), int(nodePositions[nodeNumber,1] + halfWindowSize[1]+1) ),
                                slice( int(nodePositions[nodeNumber,2] - halfWindowSize[2]), int(nodePositions[nodeNumber,2] + halfWindowSize[2]+1) ) ]

                imagette2 = im2[subVolSlice2].copy()

                # Extract initial F for correlation, remove int() part of displacement since it's already used to extract imagette2
                Finit = Ffield[nodeNumber].copy()
                Finit[0:3,-1] -= nodeDisplacement.copy()
                #print "Finit:\n", Finit, '\n'
                #Finit = numpy.eye(4)

                if ( numpy.array( imagette1.shape ) - numpy.array( imagette2.shape ) == margin*2 ).all():
293
294
295
296
297
298
299
300
                    returnStatus = 2

                # Failed innermost condition -- im1 and im2 margin alignment -- this is a harsh condition
                else:
                    returnStatus = -4
                    imagette1 = None
                    imagette2 = None
                    
Edward Andò's avatar
Edward Andò committed
301

302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
            # Failed mask or greylevel condition
            else:
                returnStatus = -5
                imagette1 = None
                imagette2 = None
        
        # Failed non-NaN components in F
        else:
            returnStatus = -6
            imagette1 = None
            imagette2 = None
            
        return { 'imagette1': imagette1, 'imagette2': imagette2, 'imagette1mask': imagette1mask, 'returnStatus': returnStatus, 'Finit': Finit, 'nodeDisplacement': nodeDisplacement }
            
        
    
    if mpi:
        import mpi4py.MPI

        mpiComm = mpi4py.MPI.COMM_WORLD
        mpiSize = mpiComm.Get_size()
        mpiRank = mpiComm.Get_rank()
        mpiStatus = mpi4py.MPI.Status()

        boss = mpiSize-1
        
        numberOfWorkers = mpiSize - 1
        workersActive   = numpy.zeros( numberOfWorkers )
    else:
        numberOfWorkers = 1
        workersActive = numpy.array( [ 0 ] )

    # Check input sanity
    if type( halfWindowSize ) == int or type( halfWindowSize ) == float:
        halfWindowSize = [ halfWindowSize ]*3

    # Check minSubVolume
    if minSubVolume is None:
        minSubVolume = int( ((1+min(halfWindowSize)*2)**3) * 0.5 )
    else:
        minSubVolume = int( ((1+min(halfWindowSize)*2)**3) * minSubVolume )

    # Check F field
    if Ffield is None:
        Ffield = numpy.zeros( ( numberOfNodes, 4, 4 ) )
        for nodeNumber in xrange( numberOfNodes ): Ffield[nodeNumber] = numpy.eye( 4 )
        
    # Add nodes to a queue -- mostly useful for MPI
    q = Queue.Queue() 
    for node in range( numberOfNodes ): q.put( node )
    finishedNodes   = 0

    subpixelError        = numpy.zeros( ( numberOfNodes ) )
    subpixelIterations   = numpy.zeros( ( numberOfNodes ) )
    subpixelReturnStatus = numpy.zeros( ( numberOfNodes ) )
    subpixelDeltaFnorm   = numpy.zeros( ( numberOfNodes ) )

    writeReturns = False

    #print( "\n\n\tStarting Node correlation 2 (subpixel)" )
    while finishedNodes != numberOfNodes:
        # If there are workers not working, satify their requests...
        #   Note: this condition is alyas true if we are not in MPI and there are jobs to do
        if workersActive.sum() < numberOfWorkers and not q.empty():
            worker  = numpy.where( workersActive == False )[0][0]
            # Get the next node off the queue
            nodeNumber    = q.get()
            
            imagetteReturns = getImagettes( nodeNumber, im1, im2, Ffield, nodePositions, margin, halfWindowSize, greyLowThreshold, im1mask)
            
            if imagetteReturns[ 'returnStatus' ] == 2:
                if mpi:
                    # build message for lukas kanade worker
                    m = {   'nodeNumber'  : nodeNumber,
                            'im1'         : imagetteReturns['imagette1'],
                            'im2'         : imagetteReturns['imagette2'],
                            'im1mask'     : imagetteReturns['imagette1mask'],
                            'Finit'       : imagetteReturns['Finit'],
                            #'FinitBinRatio' : FinitBinRatio,
                            'margin'        : margin,
                            'maxIterations' : maxIterations,
                            'minFchange'  : minFchange,
                            'interpolationOrder' : interpolationOrder,
                            'interpolator'  : interpolator,
                            'nodeDisplacement' : imagetteReturns['nodeDisplacement']
                        }
                
                    #print "\tBoss: sending node {} to worker {}".format( nodeNumber, worker )
                    mpiComm.send( m, dest=worker, tag=1 )
                    
                    # Mark this worker as working
                    workersActive[ worker ] = True
                    
                    # NOTE: writeReturns is defined later when receiving messages
                    
               
                else: # Not MPI
                    returns = spam.DIC.correlate.lucasKanade(   imagetteReturns['imagette1'],
                                                                imagetteReturns['imagette2'],
                                                                im1mask            = imagetteReturns['imagette1mask'],
                                                                Finit              = imagetteReturns['Finit'],
Edward Andò's avatar
Edward Andò committed
403
404
405
406
407
408
409
410
411
                                                                # TODO: Issue 61 -- this might be better placed in the subVolSlice2
                                                                # 2018-01-15 EA and OS: We have im2 bigger than im1 so there is automatically a margin
                                                                # Need at least a margin for the gradient calculation
                                                                margin             = 1, # see top of this file for compensation
                                                                maxIterations      = maxIterations,
                                                                minFchange         = minFchange,
                                                                interpolationOrder = interpolationOrder,
                                                                verbose            = False,
                                                                imShowProgress     = False )
412
413
414
415
416
417
                    nodeDisplacement = imagetteReturns['nodeDisplacement']
                    writeReturns = True
                    
            else: # Regardless of MPI or single proc
                # Failed to extract imagettes or something
                subpixelError[        nodeNumber ] = numpy.inf
Edward Andò's avatar
Edward Andò committed
418
                subpixelIterations[   nodeNumber ] = 0
419
                subpixelReturnStatus[ nodeNumber ] = imagetteReturns['returnStatus']
Edward Andò's avatar
Edward Andò committed
420
                subpixelDeltaFnorm[   nodeNumber ] = numpy.inf
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
                finishedNodes                     += 1

        # Otherwise spend time looking waiting for replies from workers
        elif mpi:
            message = mpiComm.recv( source=mpi4py.MPI.ANY_SOURCE, tag=2, status=mpiStatus)
            tag     = mpiStatus.Get_tag()
            if tag == 2:
                worker          = message[0]
                nodeNumber      = message[1]
                returns         = message[2]
                nodeDisplacement= message[3]
                #print "\tBoss: received node {} from worker {}".format( nodeNumber, worker )
                workersActive[worker]   = False
                writeReturns            = True                
            else:
                print "\tBoss: Don't recognise tag ", tag
Edward Andò's avatar
Edward Andò committed
437
        
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
        # If we have new DVC returns, save them in our output matrices
        if writeReturns:
                finishedNodes          += 1 
                writeReturns = False
                # Overwrite transformation operator for this node
                Ffield[ nodeNumber ] = returns['Fcentre']
                # Add back in the translation from the initial guess
                Ffield[ nodeNumber, 0:3, 3 ] += nodeDisplacement

                subpixelError[        nodeNumber ] = returns['error']
                subpixelIterations[   nodeNumber ] = returns['iterationNumber']
                subpixelReturnStatus[ nodeNumber ] = returns['returnStatus']
                subpixelDeltaFnorm[   nodeNumber ] = returns['deltaFnorm']
                print( "\r\t\tCorrelating node {:04d} of {:04d}".format( nodeNumber+1, numberOfNodes ) ),
                print( "error={:05.0f}\titerations={:02d}\tdeltaFnorm={:0.5f}\treturnStat={}".format( returns['error'], returns['iterationNumber'], returns['deltaFnorm'], returns['returnStatus'] ) ),



    # tidy up, send message type 3 to all workers
    if mpi:
        for worker in range( numberOfWorkers ): mpiComm.send( None, dest=worker, tag=3 )
459
460
461
462
463

    return {    "Ffield": Ffield,
                "subpixelError": subpixelError,
                "subpixelIterations": subpixelIterations,
                "subpixelReturnStatus": subpixelReturnStatus,
Edward Andò's avatar
Edward Andò committed
464
                "subpixelDeltaFnorm": subpixelDeltaFnorm,
465
             }