From bfaf2bf3acab28973eb5db2b6a75ff609f9d2bd2 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 24 Apr 2019 11:29:49 -0400 Subject: [PATCH] Add a tool for loading dark images from cxi files --- CDTools/tools/data.py | 53 ++++++++++++++++++++++++++++++++++++++++ tests/conftest.py | 22 +++++++++++++---- tests/tools/test_data.py | 25 ++++++++++++++++++- 3 files changed, 94 insertions(+), 6 deletions(-) diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index 4d1ad9c..96a0395 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -13,6 +13,7 @@ __all__ = ['get_entry_info', 'get_wavelength', 'get_detector_geometry', 'get_mask', + 'get_dark', 'get_data', 'get_ptycho_translations', 'create_cxi', @@ -21,6 +22,7 @@ __all__ = ['get_entry_info', 'add_source', 'add_detector', 'add_mask', + 'add_dark', 'add_data', 'add_ptycho_translations'] @@ -284,6 +286,35 @@ def get_mask(cxi_file): return None +def get_dark(cxi_file): + """Returns an array with a dark image to use for initialization of a background model + + This looks for a set of dark images at + entry_1/instrument_1/detector_1/data_dark. If the darks exist, it will + return the mean of the array along all axes but the last two. That is, + if the dark image is a single image, it will return that image. If it + is a stack of images, it will return the mean along the stack axis. + + If the darks do not exist, it will return None + + Args: + cxi_file (h5py.File) : a file object to be read + + Returns: + np.array : An array storing the dark image + """ + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1/data_dark' in i1: + darks = np.array(i1['detector_1/data_dark']) + dims = tuple(range(len(darks.shape) - 2)) + darks = np.nanmean(darks,axis=dims) + else: + darks = None + + return darks + + + def get_data(cxi_file, cut_zeroes = True): """Returns an array with the full stack of detector data defined in the cxi file object @@ -508,6 +539,28 @@ def add_mask(cxi_file, mask): d1.create_dataset('mask',data=mask_to_save) +def add_dark(cxi_file, dark): + """Adds the specified dark image to a cxi file + + It places the dark image data into the data_dark dataset under + entry_1/instrument_1/detector_1. + + Args: + cxi_file (h5py.File) : The file to add the mask to + dark (array_like) : The dark image(s) to save out to the file + """ + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1' not in i1: + i1.create_group('detector_1') + d1 = i1['detector_1'] + if isinstance(dark, t.Tensor): + dark = dark.detach().cpu().numpy() + + d1.create_dataset('data_dark',data=dark) + + def add_data(cxi_file, data, axes=None): """Adds the specified data to the cxi file diff --git a/tests/conftest.py b/tests/conftest.py index 6e90f24..a1b6a0c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -103,10 +103,15 @@ def ptycho_cxi_1(): # Remember the format for the CXI file differs from the format used # internally - mask = np.zeros((100,256,256)).astype(np.uint32) - expected['mask'] = np.ones((100,256,256)).astype(np.uint8) + mask = np.zeros((256,256)).astype(np.uint32) + expected['mask'] = np.ones((256,256)).astype(np.uint8) d1f.create_dataset('mask',data=mask) + # Create an initial background + dark = np.ones((256,256)) * 0.01 + expected['dark'] = dark + d1f.create_dataset('data_dark', data=dark) + data1f = e1f.create_group('data_1') data = np.random.rand(100,256,256).astype(np.float32) @@ -187,6 +192,12 @@ def ptycho_cxi_2(): # internally expected['mask'] = None + # Test with a set of dark images + dark = np.ones((10,256,256)) * 0.01 + expected['dark'] = np.nanmean(dark,axis=0) + d1f.create_dataset('data_dark', data=dark) + + data1f = e1f.create_group('data_1') data = np.random.rand(100,256,256).astype(np.float32) @@ -257,10 +268,11 @@ def ptycho_cxi_3(): # Remember the format for the CXI file differs from the format used # internally - mask = np.ones((100,256,256)).astype(np.uint32) * 0x00001000 - expected['mask'] = np.ones((100,256,256)).astype(np.uint8) + mask = np.ones((256,256)).astype(np.uint32) * 0x00001000 + expected['mask'] = np.ones((256,256)).astype(np.uint8) d1f.create_dataset('mask',data=mask) - + expected['dark'] = None + data1f = e1f.create_group('data_1') data = np.random.rand(100,256,256).astype(np.float32) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 573a8e5..5045365 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -64,7 +64,16 @@ def test_get_mask(test_ptycho_cxis): continue assert np.all(data.get_mask(cxi) == expected['mask']) - + +def test_get_dark(test_ptycho_cxis): + for cxi, expected in test_ptycho_cxis: + dark = data.get_dark(cxi) + if dark is None: + assert expected['dark'] is None + else: + assert np.allclose(dark, expected['dark']) + + def test_get_data(test_ptycho_cxis): for cxi, expected in test_ptycho_cxis: patterns, axes = data.get_data(cxi) @@ -188,7 +197,21 @@ def test_add_mask(tmp_path): read_mask = data.get_mask(f) assert np.all(mask == read_mask) + +def test_add_dark(tmp_path): + dark = np.random.rand(350,620) + + with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f: + data.add_dark(f, dark) + + with h5py.File(tmp_path / 'test_add_dark.cxi') as f: + read_dark = data.get_dark(f) + + print(dark.shape) + assert np.allclose(dark, read_dark) + + def test_add_data(tmp_path): # First test from numpy, with axes