image_processing.py 52.2 KB
Newer Older
1
2
3
4
5
"""
Created on Jun 16, 2016

@author: Chris Smith -- csmith55@utk.edu
"""
6

7
from __future__ import division, print_function, absolute_import
8
import os
9
from multiprocessing import cpu_count
10
from warnings import warn
11
12
import matplotlib.pyplot as plt
import numpy as np
13
from scipy.optimize import leastsq
Chris Smith's avatar
Chris Smith committed
14
from scipy.signal import blackman
15
from sklearn.utils import gen_batches
16

Chris Smith's avatar
Chris Smith committed
17
18
from ..io.hdf_utils import getH5DsetRefs, copyAttributes, linkRefs, findH5group, calc_chunks, link_as_main, \
    check_for_old
19
20
21
from ..io.io_hdf5 import ioHDF5
from ..io.io_utils import getAvailableMem
from ..io.microdata import MicroDataGroup, MicroDataset
Somnath, Suhas's avatar
Somnath, Suhas committed
22
from ..io.translators.utils import get_position_slicing, make_position_mat, get_spectral_slicing
Chris Smith's avatar
Chris Smith committed
23
from .svd_utils import _get_component_slice
24

Chris Smith's avatar
Chris Smith committed
25
26
27
28
29
30
31
32
33
34
35
36
# windata32 = np.dtype([('Image Data', np.float32)])
# absfft32 = np.dtype([('FFT Magnitude', np.float32)])
# winabsfft32 = np.dtype([('Image Data', np.float32), ('FFT Magnitude', np.float32)])
# wincompfft32 = np.dtype([('Image Data', np.float32), ('FFT Real', np.float32), ('FFT Imag', np.float32)])
windata32 = np.dtype({'names': ['Image Data'],
                      'formats': [np.float32]})
absfft32 = np.dtype({'names': ['FFT Magnitude'],
                     'formats': [np.float32]})
winabsfft32 = np.dtype({'names': ['Image Data', 'FFT Magnitude'],
                        'formats': [np.float32, np.float32]})
wincompfft32 = np.dtype({'names': ['Image Data', 'FFT Real', 'FFT Imag'],
                         'formats': [np.float32, np.float32, np.float32]})
37

Unknown's avatar
Unknown committed
38

39
40
41
42
43
44
class ImageWindow(object):
    """
    This class will handle the reading of a raw image file, creating windows from it, and writing those
    windows to an HDF5 file.
    """

45
    def __init__(self, h5_main, max_RAM_mb=1024, cores=None, reset=True, **image_args):
46
47
48
        """
        Setup the image windowing

49
50
        Parameters
        ----------
Chris Smith's avatar
Chris Smith committed
51
52
53
54
55
56
57
58
59
60
        h5_main : h5py.Dataset
            HDF5 Dataset containing the data to be windowed
        max_RAM_mb : int, optional
            integer maximum amount of ram, in Mb, to use in windowing
            Default 1024
        cores : int, optional
            integer number of [logical] CPU cores to use in windowing
            Defualt None, use number of available cores minus 2
        reset : Boolean, optional
            should all data in the hdf5 file be deleted
61

62
        """
63
64
        self.h5_file = h5_main.file
        self.hdf = ioHDF5(h5_main.file)
Unknown's avatar
Unknown committed
65

66
        # Ensuring that at least one core is available for use / 2 cores are available for other use
Unknown's avatar
Unknown committed
67
68
69
        max_cores = max(1, cpu_count() - 2)
        #         print 'max_cores',max_cores
        if cores is not None:
70
71
            cores = min(round(abs(cores)), max_cores)
        else:
Unknown's avatar
Unknown committed
72
73
            cores = max_cores
        self.cores = int(cores)
74

Unknown's avatar
Unknown committed
75
        self.max_memory = min(max_RAM_mb * 1024 ** 2, 0.75 * getAvailableMem())
76
        if self.cores != 1:
Unknown's avatar
Unknown committed
77
            self.max_memory = int(self.max_memory / 2)
Chris Smith's avatar
Chris Smith committed
78

Chris Smith's avatar
Chris Smith committed
79
80
81
        '''
        Initialize class variables to None
        '''
82
        self.h5_raw = h5_main
83
84
85
        self.h5_norm = None
        self.h5_wins = None
        self.h5_clean = None
Chris Smith's avatar
Chris Smith committed
86
87
88
        self.h5_noise = None
        self.h5_fft_clean = None
        self.h5_fft_noise = None
89

90
    def do_windowing(self, win_x=None, win_y=None, win_step_x=1, win_step_y=1,
Chris Smith's avatar
Chris Smith committed
91
                     win_fft=None, *args, **kwargs):
92
93
94
        """
        Extract the windows from the normalized image and write them to the file

95
96
        Parameters
        ----------
Chris Smith's avatar
Chris Smith committed
97
98
99
100
101
102
103
104
105
106
107
108
        win_x : int, optional
            size of the window, in pixels, in the horizontal direction
            Default None, a guess will be made based on the FFT of the image
        win_y : int, optional
            size of the window, in pixels, in the vertical direction
            Default None, a guess will be made based on the FFT of the image
        win_step_x : int, optional
            step size, in pixels, to take between windows in the horizontal direction
            Default 1
        win_step_y : int, optional
            step size, in pixels, to take between windows in the vertical direction
            Default 1
Unknown's avatar
Unknown committed
109
        win_fft : str, optional
Chris Smith's avatar
Chris Smith committed
110
111
112
113
            What kind of fft should be stored with the windows.  Options are
            None - Only the window
            'abs' - Only the magnitude of the fft
            'data+abs' - The window and magnitude of the fft
Chris Smith's avatar
Chris Smith committed
114
            'data+complex' - The window and the complex fft
Chris Smith's avatar
Chris Smith committed
115
            Default None
116
117
118

        Returns
        -------
Chris Smith's avatar
Chris Smith committed
119
120
121
        h5_wins : HDF5 Dataset
            Dataset containing the flattened windows

122
        """
123
        h5_main = self.h5_raw
124

Chris Smith's avatar
Chris Smith committed
125
        if win_fft == 'data' or win_fft is None:
Chris Smith's avatar
Chris Smith committed
126
            win_fft = 'data'
Chris Smith's avatar
Chris Smith committed
127
            win_type = windata32
Chris Smith's avatar
Chris Smith committed
128
            win_func = self.win_data_func
Chris Smith's avatar
Chris Smith committed
129
        elif win_fft == 'abs':
Chris Smith's avatar
Chris Smith committed
130
131
            win_type = absfft32
            win_func = self.abs_fft_func
Chris Smith's avatar
Chris Smith committed
132
        elif win_fft == 'data+abs':
Chris Smith's avatar
Chris Smith committed
133
            win_type = winabsfft32
Chris Smith's avatar
Chris Smith committed
134
            win_func = self.win_abs_fft_func
Chris Smith's avatar
Chris Smith committed
135
        elif win_fft == 'data+complex':
Chris Smith's avatar
Chris Smith committed
136
            win_type = wincompfft32
Chris Smith's avatar
Chris Smith committed
137
            win_func = self.win_comp_fft_func
Chris Smith's avatar
Chris Smith committed
138
139
140
141
142
        else:
            warn('Invalid FFT option supplied.  Windowing will default to data only.')
            win_fft = 'data'
            win_type = windata32
            win_func = self.win_data_func
Chris Smith's avatar
Chris Smith committed
143

144
145
        '''
        If a window size has not been specified, obtain a guess value from 
Chris Smith's avatar
Chris Smith committed
146
        window_size_extract
147
        '''
Chris Smith's avatar
Chris Smith committed
148
149
150
151
152
        win_test, psf_width = self.window_size_extract(*args, **kwargs)
        if win_x is None:
            win_x = win_test
        if win_y is None:
            win_y = win_test
Chris Smith's avatar
Chris Smith committed
153

Chris Smith's avatar
Chris Smith committed
154
155
        image, h5_wins, win_pos_mat, have_old = self._setup_window_h5(h5_main, psf_width, win_fft, win_step_x,
                                                                      win_step_y, win_type, win_x, win_y)
Chris Smith's avatar
Chris Smith committed
156

Chris Smith's avatar
Chris Smith committed
157
158
159
        if have_old:
            self.h5_wins = h5_wins
            return h5_wins
Chris Smith's avatar
Chris Smith committed
160

Chris Smith's avatar
Chris Smith committed
161
162
        n_wins = win_pos_mat.shape[0]
        win_pix = win_x * win_y
163
164
165
166

        '''
        Create slice object from the positions
        '''
Unknown's avatar
Unknown committed
167
        win_slices = [[slice(x, x + win_x), slice(y, y + win_y)] for x, y in win_pos_mat]
168

169
        '''
170
        Calculate the size of a given batch that will fit in the available memory
171
        '''
Unknown's avatar
Unknown committed
172
        mem_per_win = win_x * win_y * h5_wins.dtype.itemsize
173
        if self.cores is None:
Unknown's avatar
Unknown committed
174
            free_mem = self.max_memory - image.size * image.itemsize
175
        else:
Unknown's avatar
Unknown committed
176
177
            free_mem = self.max_memory * 2 - image.size * image.itemsize
        batch_size = int(free_mem / mem_per_win)
178
        batch_slices = gen_batches(n_wins, batch_size)
179

180
        for ibatch, batch in enumerate(batch_slices):
Unknown's avatar
Unknown committed
181
            batch_wins = np.zeros([batch.stop - batch.start, win_pix], dtype=win_type)
182
183
184
185
            '''
            Read each slice and write it to the dataset
            '''
            for islice, this_slice in enumerate(win_slices[batch]):
Unknown's avatar
Unknown committed
186
187
                iwin = ibatch * batch_size + islice
                selected = iwin % np.rint(n_wins / 10) == 0
188
189

                if selected:
Unknown's avatar
Unknown committed
190
                    per_done = np.rint(100 * iwin / n_wins)
191
192
193
194
195
196
197
                    print('Windowing Image...{}% --pixels {}-{}, step # {}'.format(per_done,
                                                                                   (this_slice[0].start,
                                                                                    this_slice[1].start),
                                                                                   (this_slice[0].stop,
                                                                                    this_slice[1].stop),
                                                                                   islice))

Chris Smith's avatar
Chris Smith committed
198
                batch_wins[islice] = win_func(image[this_slice]).flatten()
199
200
201

            h5_wins[batch] = batch_wins
            self.hdf.flush()
Unknown's avatar
Unknown committed
202

203
        self.h5_wins = h5_wins
Unknown's avatar
Unknown committed
204

205
206
        return h5_wins

Chris Smith's avatar
Chris Smith committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
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
293
294
295
296
297
298
299
300
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
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
    @staticmethod
    def _check_win_parameters(h5_main, win_step_x, win_step_y, win_x, win_y):
        """
        Check the window parameters
        
        Parameters
        ----------
        h5_main : h5py.Dataset
            Dataset containing the Raw Image
        win_step_x : uint
            Step size in the x-direction between windows.
        win_step_y : uint
            Step size in the y-direction between windows.
        win_x : uint
            Size of the window in the x-direction.
        win_y : uint
            Size of the window in the y-direction.

        Returns
        -------
        image : numpy.ndarray
            Array containing the original image reshaped by position.
        win_step_x : uint
            Step size in the x-direction between windows.
        win_step_y : uint
            Step size in the y-direction between windows.
        win_x : uint
            Size of the window in the x-direction.
        win_y : uint
            Size of the window in the y-direction.
        
        """
        '''
        Get the position indices of h5_main and reshape the flattened image back
        '''
        try:
            h5_pos = h5_main.file[h5_main.attrs['Position_Indices']][()]
            x_pix = len(np.unique(h5_pos[:, 0]))
            y_pix = len(np.unique(h5_pos[:, 1]))

        except KeyError:
            '''
            Position Indices dataset does not exist
            Assume square image
            '''
            x_pix = np.int(np.sqrt(h5_main.size))
            y_pix = x_pix

        except:
            raise
        image = h5_main[()].reshape(x_pix, y_pix)
        '''
            Step size must be less than 1/4th the image size
            '''
        win_step_x = min(x_pix / 4, win_step_x)
        win_step_y = min(y_pix / 4, win_step_y)
        '''
            Prevent windows from being less that twice the step size and more than half the image size
            '''
        win_x = max(2 * win_step_x, min(x_pix, win_x))
        win_y = max(2 * win_step_y, min(y_pix, win_y))
        print('Optimal window size determined to be {wx}x{wy} pixels.'.format(wx=win_x, wy=win_y))

        return image, win_step_x, win_step_y, win_x, win_y

    def _setup_window_h5(self, h5_main, psf_width, win_fft, win_step_x, win_step_y, win_type, win_x,
                         win_y):
        """
        Setup the hdf5 group for the windows
        
        Parameters
        ----------
        h5_main : h5py.Dataset
            Dataset containing the Raw Image
        psf_width : uint
            psf_width???  Someone who knows what this is should fill it in 
        win_fft : 
        win_step_x : uint
            Step size in the x-direction between windows.
        win_step_y : uint
            Step size in the y-direction between windows.
        win_type
        win_x : uint
            Size of the window in the x-direction.
        win_y : uint
            Size of the window in the y-direction.

        Returns
        -------
        

        """

        image, win_step_x, win_step_y, win_x, win_y = self._check_win_parameters(h5_main,
                                                                                 win_step_x, win_step_y,
                                                                                 win_x, win_y)

        '''
        Build the Spectroscopic and Position Datasets 
        '''
        ds_pix_inds, ds_pix_vals, ds_pos_inds, ds_pos_vals, win_pos_mat = self._get_window_pos_spec(
            image, win_step_x, win_step_y, win_x, win_y)
        im_x, im_y = image.shape
        n_wins = win_pos_mat.shape[0]
        win_pix = win_x * win_y

        '''
        Calculate the chunk size
        '''
        win_chunks = calc_chunks([n_wins, win_pix], win_type.itemsize, unit_chunks=[1, win_pix])

        parent = h5_main.parent

        check_parameters = {'fft_mode': win_fft,
                            'psf_width': psf_width,
                            'win_x': win_x,
                            'win_y': win_y,
                            'win_step_x': win_step_x,
                            'win_step_y': win_step_y,
                            'image_x': im_x,
                            'image_y': im_y}
        basename = h5_main.name.split('/')[-1]

        old_group = check_for_old(h5_main, '-Windowing', check_parameters)

        if old_group is not None:
            old = True
            h5_wins = old_group['Image_Windows']

        else:
            old = False
            '''
            Create the Windows Dataset and Datagroup
            '''
            ds_windows = MicroDataset('Image_Windows',
                                      data=[],
                                      maxshape=[n_wins, win_pix],
                                      dtype=win_type,
                                      chunking=win_chunks,
                                      compression='gzip')

            ds_group = MicroDataGroup(basename + '-Windowing_', parent.name[1:])
            ds_group.addChildren([ds_windows, ds_pos_inds, ds_pix_inds,
                                  ds_pos_vals, ds_pix_vals])
            ds_group.attrs['win_x'] = win_x
            ds_group.attrs['win_y'] = win_y
            ds_group.attrs['win_step_x'] = win_step_x
            ds_group.attrs['win_step_y'] = win_step_y
            ds_group.attrs['image_x'] = im_x
            ds_group.attrs['image_y'] = im_y
            ds_group.attrs['psf_width'] = psf_width
            ds_group.attrs['fft_mode'] = win_fft
            image_refs = self.hdf.writeData(ds_group)

            '''
            Get the hdf5 objects for the windows and ancillary datasets
            '''
            h5_wins = getH5DsetRefs(['Image_Windows'], image_refs)[0]

            '''
            Link references to windowed dataset
            '''
            aux_ds_names = ['Position_Indices', 'Position_Values', 'Spectroscopic_Indices', 'Spectroscopic_Values']
            linkRefs(h5_wins, getH5DsetRefs(aux_ds_names, image_refs))

            self.hdf.flush()

        return image, h5_wins, win_pos_mat, old

    @staticmethod
    def _get_window_pos_spec(image, win_step_x, win_step_y, win_x, win_y):
        """
        Create the position and spectroscopic datasets for the windows.
        
        Parameters
        ----------
        image : numpy.ndarray
            Raw Image
        win_step_x : uint
            Step size in the x-direction between windows.
        win_step_y : uint
            Step size in the y-direction between windows.
        win_x : uint
            Size of the window in the x-direction.
        win_y : uint
            Size of the window in the y-direction.

        Returns
        -------
        ds_pix_inds : MicroDataset
            Spectroscopic Indices of the windows
        ds_pix_vals : MicroDataset
            Spectroscopic Values of the windows
        ds_pos_inds : MicroDataset
            Position Indices of the windows
        ds_pos_vals : MicroDataset
            Position Values of the windows
        win_pos_mat : numpy.ndarray
            Array containing the positions of the window origins

        """
        im_x, im_y = image.shape
        x_steps = np.arange(0, im_x - win_x + 1, win_step_x, dtype=np.uint32)
        y_steps = np.arange(0, im_y - win_y + 1, win_step_y, dtype=np.uint32)
        nx = len(x_steps)
        ny = len(y_steps)
        win_pos_mat = np.array([np.repeat(x_steps, ny), np.tile(y_steps, nx)], dtype=np.uint32).T
        win_pix_mat = make_position_mat([win_x, win_y]).T

        '''
        Set up the HDF5 Group and Datasets for the windowed data
        '''
        ds_pos_inds = MicroDataset('Position_Indices', data=win_pos_mat, dtype=np.int32)
        ds_pix_inds = MicroDataset('Spectroscopic_Indices', data=win_pix_mat, dtype=np.int32)
        ds_pos_vals = MicroDataset('Position_Values', data=win_pos_mat, dtype=np.float32)
        ds_pix_vals = MicroDataset('Spectroscopic_Values', data=win_pix_mat, dtype=np.float32)
        pos_labels = get_position_slicing(['Window Origin X', 'Window Origin Y'])
        ds_pos_inds.attrs['labels'] = pos_labels
        ds_pos_inds.attrs['units'] = ['pixel', 'pixel']
        ds_pos_vals.attrs['labels'] = pos_labels
        ds_pos_vals.attrs['units'] = ['pixel', 'pixel']
        pix_labels = get_spectral_slicing(['U', 'V'])
        ds_pix_inds.attrs['labels'] = pix_labels
        ds_pix_inds.attrs['units'] = ['pixel', 'pixel']
        ds_pix_vals.attrs['labels'] = pix_labels
        ds_pix_vals.attrs['units'] = ['pixel', 'pixel']

        return ds_pix_inds, ds_pix_vals, ds_pos_inds, ds_pos_vals, win_pos_mat

Chris Smith's avatar
Chris Smith committed
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
    @staticmethod
    def win_data_func(image):
        """
        Returns the input image in the `windata32` format

        Parameters
        ----------
        image : numpy.ndarray
            Windowed image to take the FFT of

        Returns
        -------
        windows : numpy.ndarray
            Array the image in the windata32 format
        """
        windows = np.empty_like(image, dtype=windata32)
        windows['Image Data'] = image

        return windows

Chris Smith's avatar
Chris Smith committed
456
457
458
    @staticmethod
    def abs_fft_func(image):
        """
Chris Smith's avatar
Chris Smith committed
459
        Take the 2d FFT of each window in `windows` and return in the proper form.
Chris Smith's avatar
Chris Smith committed
460

Chris Smith's avatar
Chris Smith committed
461
462
463
464
465
466
467
468
469
470
        Parameters
        ----------
        image : numpy.ndarray
            Windowed image to take the FFT of

        Returns
        -------
        windows : numpy.ndarray
            Array of the Magnitude of the FFT of each window for the input
            `image`
Chris Smith's avatar
Chris Smith committed
471

Chris Smith's avatar
Chris Smith committed
472
473
474
475
476
477
478
479
        """
        windows = np.empty_like(image, dtype=absfft32)
        windows['FFT Magnitude'] = np.abs(np.fft.fftshift(np.fft.fft2(image)))

        return windows

    @staticmethod
    def win_abs_fft_func(image):
Chris Smith's avatar
Chris Smith committed
480
481
482
        """
        Take the 2d FFT of each window in `windows` and return in the proper form.

Chris Smith's avatar
Chris Smith committed
483
484
485
486
487
488
489
490
491
492
493
        Parameters
        ----------
        image : numpy.ndarray
            Windowed image to take the FFT of

        Returns
        -------
        windows : numpy.ndarray
            Array of windows and the Magnitude of the FFT of each window for the input
            `image`

Chris Smith's avatar
Chris Smith committed
494
        """
Chris Smith's avatar
Chris Smith committed
495
496
497
        windows = np.empty_like(image, dtype=winabsfft32)
        windows['Image Data'] = image
        windows['FFT Magnitude'] = np.abs(np.fft.fftshift(np.fft.fft2(image)))
Chris Smith's avatar
Chris Smith committed
498
499
500

        return windows

Chris Smith's avatar
Chris Smith committed
501
502
    @staticmethod
    def win_comp_fft_func(image):
Chris Smith's avatar
Chris Smith committed
503
504
505
        """
        Take the 2d FFT of each window in `windows` and return in the proper form.

Chris Smith's avatar
Chris Smith committed
506
507
508
509
510
511
512
513
514
515
        Parameters
        ----------
        image : numpy.ndarray
            Windowed image to take the FFT of

        Returns
        -------
        windows : numpy.ndarray
            Array of windows and the FFT of each window for the input `image`

Chris Smith's avatar
Chris Smith committed
516
        """
Chris Smith's avatar
Chris Smith committed
517
518
519
        windows = np.empty_like(image, dtype=wincompfft32)
        windows['Image Data'] = image
        win_fft = np.fft.fftshift(np.fft.fft2(image))
Chris Smith's avatar
Chris Smith committed
520
521
522
523
524
        windows['FFT Real'] = win_fft.real
        windows['FFT Imag'] = win_fft.imag

        return windows

525
526
527
528
    def build_clean_image(self, h5_win=None):
        """
        Reconstructs the cleaned image from the windowed dataset

529
        Parameters
530
        ----------
531
532
        h5_win : HDF5 dataset , optional
            The windowed image to be reconstructed.
533

534
535
536
537
        Returns
        -------
        h5_clean : HDF5 dataset
            The cleaned image
Chris Smith's avatar
Chris Smith committed
538

539
540
541
542
543
544
        """
        if h5_win is None:
            if self.clean_wins is None:
                warn('You must clean the image before rebuilding it.')
                return
            h5_win = self.clean_wins
Unknown's avatar
Unknown committed
545

546
547
548
549
550
551
552
553
554
        '''
        Get basic windowing information from attributes of 
        h5_win
        '''
        im_x = h5_win.parent.attrs['image_x']
        im_y = h5_win.parent.attrs['image_y']
        win_x = h5_win.parent.attrs['win_x']
        win_y = h5_win.parent.attrs['win_y']
        win_step_x = h5_win.parent.attrs['win_step_x']
Unknown's avatar
Unknown committed
555
556
        win_step_y = h5_win.parent.attrs['win_step_x']

557
558
559
        '''
        Calculate the steps taken to create original windows
        '''
Unknown's avatar
Unknown committed
560
561
562
        x_steps = np.arange(0, im_x - win_x + 1, win_step_x)
        y_steps = np.arange(0, im_y - win_y + 1, win_step_y)

563
564
565
566
567
        '''
        Initialize arrays to hold summed windows and counts for each position
        '''
        counts = np.zeros([im_x, im_y], np.uint8)
        accum = np.zeros([im_x, im_y], np.float32)
Unknown's avatar
Unknown committed
568

569
570
        nx = len(x_steps)
        ny = len(y_steps)
Unknown's avatar
Unknown committed
571
572
        n_wins = nx * ny

573
574
575
        '''
        Create slice object from the positions
        '''
Unknown's avatar
Unknown committed
576
577
578
        win_slices = [[slice(x, x + win_x), slice(y, y + win_y)] for x, y in np.array([np.tile(x_steps, nx),
                                                                                       np.repeat(y_steps, ny)]).T]

579
580
581
582
583
584
        '''
        Loop over all windows.  Increment counts for window positions and 
        add current window to total.
        '''
        ones = np.ones([win_x, win_y], dtype=counts.dtype)
        for islice, this_slice in enumerate(win_slices):
Unknown's avatar
Unknown committed
585
            selected = islice % np.rint(n_wins / 10) == 0
586
            if selected:
Unknown's avatar
Unknown committed
587
                per_done = np.rint(100 * islice / n_wins)
Unknown's avatar
Unknown committed
588
589
590
591
592
593
                print('Reconstructing Image...{}% -- step # {}'.format(per_done, islice))
            counts[this_slice] += ones

            accum[this_slice] += h5_win[islice].reshape(win_x, win_y)

        clean_image = accum / counts
594
595

        clean_image[np.isnan(clean_image)] = 0
Unknown's avatar
Unknown committed
596

597
598
599
        clean_grp = MicroDataGroup('Cleaned_Image', h5_win.parent.name[1:])

        ds_clean = MicroDataset('Cleaned_Image', clean_image)
Unknown's avatar
Unknown committed
600

601
        clean_grp.addChildren([ds_clean])
Unknown's avatar
Unknown committed
602

603
604
        image_refs = self.hdf.writeData(clean_grp)
        self.hdf.flush()
Unknown's avatar
Unknown committed
605

606
        h5_clean = getH5DsetRefs(['Cleaned_Image'], image_refs)[0]
Unknown's avatar
Unknown committed
607

608
        self.h5_clean = h5_clean
Unknown's avatar
Unknown committed
609

610
611
612
613
        return h5_clean

    def clean_and_build(self, h5_win=None, components=None):
        """
614
        Rebuild the Image from the SVD results on the windows
615
616
        Optionally, only use components less than n_comp.

617
618
619
        Parameters
        ----------
        h5_win : hdf5 Dataset, optional
620
            dataset containing the windowed image which SVD was performed on
621
622
        components: {int, iterable of int, slice} optional
            Defines which components to keep
623

624
625
626
627
            Input Types
            integer : Components less than the input will be kept
            length 2 iterable of integers : Integers define start and stop of component slice to retain
            other iterable of integers or slice : Selection of component indices to retain
628

629
630
631
632
        Returns
        -------
        clean_wins : HDF5 Dataset
            the cleaned windows
Chris Smith's avatar
Chris Smith committed
633

634
635
636
637
        """

        if h5_win is None:
            if self.h5_wins is None:
638
                warn('You must perform windowing on an image followed by SVD on the window before you can clean it.')
639
640
                return
            h5_win = self.h5_wins
Chris Smith's avatar
Chris Smith committed
641
642
643
        elif 'Image Data' not in h5_win.dtype.names:
            warn('The windows must have the real space image data in them to rebuild.')
            return
644
645
646

        print('Cleaning the image by removing unwanted components.')

Chris Smith's avatar
Chris Smith committed
647
        comp_slice = _get_component_slice(components)
648
649

        '''
650
        Read the 1st n_comp components from the SVD results
651
652
653
654
655
        on h5_win
        '''
        win_name = h5_win.name.split('/')[-1]

        try:
656
            win_svd = findH5group(h5_win, 'SVD')[-1]
657

658
659
660
            h5_S = win_svd['S']
            h5_U = win_svd['U']
            h5_V = win_svd['V']
661
662

        except KeyError:
663
            warnstring = 'SVD Results for {dset} were not found in {file}.'.format(dset=win_name, file=self.image_path)
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
            warn(warnstring)
            return
        except:
            raise

        '''
        Get basic windowing information from attributes of
        h5_win
        '''
        im_x = h5_win.parent.attrs['image_x']
        im_y = h5_win.parent.attrs['image_y']
        win_x = h5_win.parent.attrs['win_x']
        win_y = h5_win.parent.attrs['win_y']
        win_step_x = h5_win.parent.attrs['win_step_x']
        win_step_y = h5_win.parent.attrs['win_step_x']

        '''
        Calculate the steps taken to create original windows
        '''
        x_steps = np.arange(0, im_x - win_x, win_step_x)
        y_steps = np.arange(0, im_y - win_y, win_step_y)

        '''
        Initialize arrays to hold summed windows and counts for each position
        '''
Chris Smith's avatar
Chris Smith committed
689
        counts = np.zeros([im_x, im_y], np.uint32)
Chris Smith's avatar
Chris Smith committed
690
        clean_image = np.zeros([im_x, im_y], np.float32)
691
692
693
694
695
696
697
698

        nx = len(x_steps)
        ny = len(y_steps)
        n_wins = nx * ny

        '''
        Create slice object from the positions
        '''
Chris Smith's avatar
Chris Smith committed
699
        h5_win_pos = h5_win.file[h5_win.attrs['Position_Indices']]
Unknown's avatar
Unknown committed
700
        win_slices = [[slice(x, x + win_x), slice(y, y + win_y)] for x, y in h5_win_pos]
701
702
703
704
705
706

        '''
        Loop over all windows.  Increment counts for window positions and
        add current window to total.
        '''
        ones = np.ones([win_x, win_y], dtype=counts.dtype)
Chris Smith's avatar
Chris Smith committed
707
        ds_V = np.dot(np.diag(h5_S[comp_slice]), h5_V['Image_Data'][comp_slice, :])
708
709

        for islice, this_slice in enumerate(win_slices):
710
            if islice % np.rint(n_wins / 10) == 0:
Unknown's avatar
Unknown committed
711
                per_done = np.rint(100 * islice / n_wins)
712
713
714
715
716
717
                print('Reconstructing Image...{}% -- step # {}'.format(per_done, islice))

            counts[this_slice] += ones

            this_win = np.dot(h5_U[islice, comp_slice], ds_V)

Chris Smith's avatar
Chris Smith committed
718
            clean_image[this_slice] += this_win.reshape(win_x, win_y)
719

Chris Smith's avatar
Chris Smith committed
720
        clean_image = np.divide(clean_image, counts)
721
722
723

        clean_image[np.isnan(clean_image)] = 0

Chris Smith's avatar
Chris Smith committed
724
725
726
        '''
        Calculate the removed noise and FFTs
        '''
Unknown's avatar
Unknown committed
727
        removed_noise = np.reshape(self.h5_raw, clean_image.shape) - clean_image
728

729
730
        fft_clean = np.fft.fft2(clean_image)
        fft_noise = np.fft.fft2(removed_noise)
731

Chris Smith's avatar
Chris Smith committed
732
733
734
        '''
        Create datasets for results, link them properly, and write them to file
        '''
Chris Smith's avatar
Chris Smith committed
735
        clean_grp = MicroDataGroup('Cleaned_Image_', win_svd.name[1:])
Chris Smith's avatar
Chris Smith committed
736
737
738
739
        ds_clean = MicroDataset('Cleaned_Image', clean_image.reshape(self.h5_raw.shape))
        ds_noise = MicroDataset('Removed_Noise', removed_noise.reshape(self.h5_raw.shape))
        ds_fft_clean = MicroDataset('FFT_Cleaned_Image', fft_clean.reshape(self.h5_raw.shape))
        ds_fft_noise = MicroDataset('FFT_Removed_Noise', fft_noise.reshape(self.h5_raw.shape))
740

Chris Smith's avatar
Chris Smith committed
741
        clean_grp.addChildren([ds_clean, ds_noise, ds_fft_clean, ds_fft_noise])
742

743
744
745
746
747
        if isinstance(comp_slice, slice):
            clean_grp.attrs['components_used'] = '{}-{}'.format(comp_slice.start, comp_slice.stop)
        else:
            clean_grp.attrs['components_used'] = comp_slice

748
749
750
751
        image_refs = self.hdf.writeData(clean_grp)
        self.hdf.flush()

        h5_clean = getH5DsetRefs(['Cleaned_Image'], image_refs)[0]
Chris Smith's avatar
Chris Smith committed
752
753
        h5_noise = getH5DsetRefs(['Removed_Noise'], image_refs)[0]
        h5_fft_clean = getH5DsetRefs(['FFT_Cleaned_Image'], image_refs)[0]
754
755
756
757
758
759
760
761
762
763
764
765
766
767
        h5_fft_noise = getH5DsetRefs(['FFT_Removed_Noise'], image_refs)[0]

        copyAttributes(self.h5_raw, h5_clean, skip_refs=False)
        copyAttributes(self.h5_raw, h5_noise, skip_refs=False)
        copyAttributes(self.h5_raw, h5_fft_clean, skip_refs=False)
        copyAttributes(self.h5_raw, h5_fft_noise, skip_refs=False)

        self.h5_clean = h5_clean
        self.h5_noise = h5_noise

        return h5_clean

    def clean_and_build_batch(self, h5_win=None, components=None):
        """
768
        Rebuild the Image from the SVD results on the windows
769
770
771
772
773
        Optionally, only use components less than n_comp.

        Parameters
        ----------
        h5_win : hdf5 Dataset, optional
774
            dataset containing the windowed image which SVD was performed on
775
776
777
778
779
780
781
782
783
784
785
786
787
        components : {int, iterable of int, slice} optional
            Defines which components to keep
            Default - None, all components kept

            Input Types
            integer : Components less than the input will be kept
            length 2 iterable of integers : Integers define start and stop of component slice to retain
            other iterable of integers or slice : Selection of component indices to retain

        Returns
        -------
        clean_wins : HDF5 Dataset
            the cleaned windows
Chris Smith's avatar
Chris Smith committed
788

789
790
791
792
        """

        if h5_win is None:
            if self.h5_wins is None:
793
                warn('You must perform windowing on an image followed by SVD on the window before you can clean it.')
794
795
                return
            h5_win = self.h5_wins
Chris Smith's avatar
Chris Smith committed
796
797
798
        elif 'Image Data' not in h5_win.dtype.names:
            warn('The windows must have the real space image data in them to rebuild.')
            return
799
800

        print('Cleaning the image by removing unwanted components.')
801

Chris Smith's avatar
Chris Smith committed
802
        comp_slice = _get_component_slice(components)
803
804

        '''
805
        Read the 1st n_comp components from the SVD results
806
807
808
809
810
        on h5_win
        '''
        win_name = h5_win.name.split('/')[-1]

        try:
811
            win_svd = findH5group(h5_win, 'SVD')[-1]
812
813
814
815
816
817

            h5_S = win_svd['S']
            h5_U = win_svd['U']
            h5_V = win_svd['V']

        except KeyError:
818
            warnstring = 'SVD Results for {dset} were not found in {file}.'.format(dset=win_name, file=self.image_path)
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
            warn(warnstring)
            return
        except:
            raise

        '''
        Get basic windowing information from attributes of
        h5_win
        '''
        im_x = h5_win.parent.attrs['image_x']
        im_y = h5_win.parent.attrs['image_y']
        win_x = h5_win.parent.attrs['win_x']
        win_y = h5_win.parent.attrs['win_y']

        '''
        Initialize arrays to hold summed windows and counts for each position
        '''
        counts = np.zeros([im_x, im_y], np.uint32)
        accum = np.zeros([im_x, im_y], np.float32)

        '''
        Create slice object from the positions
        '''
        ds_win_pos = h5_win.file[h5_win.attrs['Position_Indices']][()]
Unknown's avatar
Unknown committed
843
        win_slices = [[slice(x, x + win_x), slice(y, y + win_y)] for x, y in ds_win_pos]
844
        n_wins = ds_win_pos.shape[0]
845
846
847
848
849
        '''
        Create a matrix to add when counting.
        h5_V is usually small so go ahead and take S.V
        '''
        ones = np.ones([win_x, win_y], dtype=counts.dtype)
Chris Smith's avatar
Chris Smith committed
850
        ds_V = np.dot(np.diag(h5_S[comp_slice]), h5_V['Image Data'][comp_slice, :])
851
852
853
854

        '''
        Calculate the size of a given batch that will fit in the available memory
        '''
Unknown's avatar
Unknown committed
855
        mem_per_win = ds_V.itemsize * ds_V.shape[1]
856
        if self.cores is None:
Unknown's avatar
Unknown committed
857
            free_mem = self.max_memory - ds_V.size * ds_V.itemsize
858
        else:
Unknown's avatar
Unknown committed
859
860
            free_mem = self.max_memory * 2 - ds_V.size * ds_V.itemsize
        batch_size = int(free_mem / mem_per_win)
861
862
        batch_slices = gen_batches(n_wins, batch_size)

Chris Smith's avatar
Chris Smith committed
863
864
        print('Reconstructing in batches of {} windows.'.format(batch_size))

865
866
867
868
869
870
        '''
        Loop over all batches.  Increment counts for window positions and
        add current window to total.
        '''
        for ibatch, batch in enumerate(batch_slices):
            ds_U = h5_U[batch, comp_slice]
871
            batch_wins = np.dot(ds_U, ds_V).reshape([-1, win_x, win_y])
872
873
            del ds_U
            for islice, this_slice in enumerate(win_slices[batch]):
Unknown's avatar
Unknown committed
874
                iwin = ibatch * batch_size + islice
Chris Smith's avatar
Chris Smith committed
875
                if iwin % np.rint(n_wins / 10) == 0:
876
877
878
879
880
                    per_done = np.rint(100 * iwin / n_wins)
                    print('Reconstructing Image...{}% -- step # {}'.format(per_done, islice))

                counts[this_slice] += ones

881
                accum[this_slice] += batch_wins[islice]
882

883
        clean_image = np.divide(accum, counts)
884
885

        clean_image[np.isnan(clean_image)] = 0
886
887
888
889
890
891

        if h5_win.file.attrs['normalized']:
            '''
            Renormalize the cleaned image
            '''
            clean_image -= np.min(clean_image)
Unknown's avatar
Unknown committed
892
            clean_image = clean_image / np.max(clean_image)
893
894
895
896

        '''
        Calculate the removed noise and FFTs
        '''
Unknown's avatar
Unknown committed
897
        removed_noise = np.reshape(self.h5_raw, clean_image.shape) - clean_image
898
899
        blackman_window_rows = blackman(clean_image.shape[0])
        blackman_window_cols = blackman(clean_image.shape[1])
Unknown's avatar
Unknown committed
900
901
902
        fft_clean = np.fft.fft2(blackman_window_rows[:, np.newaxis] * clean_image * blackman_window_cols[np.newaxis, :])
        fft_noise = np.fft.fft2(
            blackman_window_rows[:, np.newaxis] * removed_noise * blackman_window_cols[np.newaxis, :])
903
904
905
906

        '''
        Create datasets for results, link them properly, and write them to file
        '''
Chris Smith's avatar
Chris Smith committed
907
        clean_grp = MicroDataGroup('Cleaned_Image_', win_svd.name[1:])
908
909
910
911
912
913
914
        ds_clean = MicroDataset('Cleaned_Image', clean_image.reshape(self.h5_raw.shape))
        ds_noise = MicroDataset('Removed_Noise', removed_noise.reshape(self.h5_raw.shape))
        ds_fft_clean = MicroDataset('FFT_Cleaned_Image', fft_clean.reshape(self.h5_raw.shape))
        ds_fft_noise = MicroDataset('FFT_Removed_Noise', fft_noise.reshape(self.h5_raw.shape))

        clean_grp.addChildren([ds_clean, ds_noise, ds_fft_clean, ds_fft_noise])

915
916
917
918
919
        if isinstance(comp_slice, slice):
            clean_grp.attrs['components_used'] = '{}-{}'.format(comp_slice.start, comp_slice.stop)
        else:
            clean_grp.attrs['components_used'] = comp_slice

920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
        image_refs = self.hdf.writeData(clean_grp)
        self.hdf.flush()

        h5_clean = getH5DsetRefs(['Cleaned_Image'], image_refs)[0]
        h5_noise = getH5DsetRefs(['Removed_Noise'], image_refs)[0]
        h5_fft_clean = getH5DsetRefs(['FFT_Cleaned_Image'], image_refs)[0]
        h5_fft_noise = getH5DsetRefs(['FFT_Removed_Noise'], image_refs)[0]

        copyAttributes(self.h5_raw, h5_clean, skip_refs=False)
        copyAttributes(self.h5_raw, h5_noise, skip_refs=False)
        copyAttributes(self.h5_raw, h5_fft_clean, skip_refs=False)
        copyAttributes(self.h5_raw, h5_fft_noise, skip_refs=False)

        self.h5_clean = h5_clean
        self.h5_noise = h5_noise

        return h5_clean

938
    def clean_and_build_separate_components(self, h5_win=None, components=None):
939
        """
940
        Rebuild the Image from the SVD results on the windows
941
942
943
944
945
        Optionally, only use components less than n_comp.

        Parameters
        ----------
        h5_win : hdf5 Dataset, optional
946
            dataset containing the windowed image which SVD was performed on
947
948
949
        components : {int, iterable of int, slice} optional
            Defines which components to keep
            Default - None, all components kept
Chris Smith's avatar
Chris Smith committed
950
            \n
951
952
953
954
            Input Types
            integer : Components less than the input will be kept
            length 2 iterable of integers : Integers define start and stop of component slice to retain
            other iterable of integers or slice : Selection of component indices to retain
955
956
957
958
959

        Returns
        -------
        clean_wins : HDF5 Dataset
            the cleaned windows
Chris Smith's avatar
Chris Smith committed
960

961
962
963
964
        """

        if h5_win is None:
            if self.h5_wins is None:
965
                warn('You must perform windowing on an image followed by SVD on the window before you can clean it.')
966
967
                return
            h5_win = self.h5_wins
Chris Smith's avatar
Chris Smith committed
968
969
970
        elif 'Image Data' not in h5_win.dtype.names:
            warn('The windows must have the real space image data in them to rebuild.')
            return
971
972

        print('Cleaning the image by removing unwanted components.')
Chris Smith's avatar
Chris Smith committed
973
        comp_slice = _get_component_slice(components)
974
975

        '''
976
        Read the 1st n_comp components from the SVD results
977
978
979
980
981
        on h5_win
        '''
        win_name = h5_win.name.split('/')[-1]

        try:
982
            win_svd = findH5group(h5_win, 'SVD')[-1]
983
984
985
986
987
988

            h5_S = win_svd['S']
            h5_U = win_svd['U']
            h5_V = win_svd['V']

        except KeyError:
989
            warnstring = 'SVD Results for {dset} were not found in {file}.'.format(dset=win_name, file=self.image_path)
990
991
992
993
994
995
996
997
998
999
1000
            warn(warnstring)
            return
        except:
            raise

        '''
        Get basic windowing information from attributes of
        h5_win
        '''
        im_x = h5_win.parent.attrs['image_x']
        im_y = h5_win.parent.attrs['image_y']