diff --git a/synpivimage/io.py b/synpivimage/io.py index d1f6485..0683288 100644 --- a/synpivimage/io.py +++ b/synpivimage/io.py @@ -1,3 +1,4 @@ +import abc import pathlib import shutil import warnings @@ -17,7 +18,19 @@ Format = Literal['json', 'json-ld'] -class Imwriter: +class Writer(abc.ABC): + """Abstract base class for writing images and metadata to disk.""" + + @abc.abstractmethod + def writeA(self, index: int, img: np.ndarray, particles: Particles = None) -> pathlib.Path: + """Write image A""" + + @abc.abstractmethod + def writeB(self, index: int, img: np.ndarray, particles: Particles = None) -> pathlib.Path: + """Write image B""" + + +class Imwriter(Writer): """Context manager for writing images and metadata to a folder. Example: @@ -148,7 +161,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): self._enabled = False -class HDF5Writer: +class HDF5Writer(Writer): """Context manager for writing images and metadata to a folder. Example: @@ -163,6 +176,8 @@ class HDF5Writer: def __init__(self, filename: Union[str, pathlib.Path], + *, + n_images: int, overwrite: bool = False, camera: Camera = None, laser: Laser = None): @@ -173,6 +188,14 @@ def __init__(self, ---------- filename : Union[str, pathlib.Path] The filename of the HDF5 file + n_images : int + Number of images to save. Needed to pre-allocate space in the HDF5 file + overwrite : bool + Overwrite the file if it exists + camera : Camera + Camera object + laser : Laser + Laser object """ self.filename = pathlib.Path(filename) @@ -180,8 +203,7 @@ def __init__(self, self.camera = camera self.laser = laser self._h5 = None - self._image_index_a = 0 - self._image_index_b = 0 + self._n_images = n_images def __enter__(self, n_imgs: Optional[int] = None): if self.filename.exists() and not self.overwrite: @@ -195,25 +217,27 @@ def __enter__(self, n_imgs: Optional[int] = None): self._h5 = h5py.File(self.filename, 'w') return self - def _get_dsimgA(self): + def _get_dsimgA(self) -> "h5py.Dataset": + """Get or create the dataset for image A""" ds_nameA = "images/img_A" if ds_nameA in self._h5: return self._h5[ds_nameA] return self._h5.create_dataset(ds_nameA, - shape=(1, self.camera.ny, self.camera.nx), - maxshape=(None, None, None), + shape=(self._n_images, self.camera.ny, self.camera.nx), + maxshape=(self._n_images, self.camera.ny, self.camera.nx), dtype='uint16') - def _get_dsimgB(self): + def _get_dsimgB(self) -> "h5py.Dataset": + """Get or create the dataset for image A""" ds_nameA = "images/img_B" if ds_nameA in self._h5: return self._h5[ds_nameA] return self._h5.create_dataset(ds_nameA, - shape=(1, self.camera.ny, self.camera.nx), - maxshape=(None, None, None), + shape=(self._n_images, self.camera.ny, self.camera.nx), + maxshape=(self._n_images, self.camera.ny, self.camera.nx), dtype='uint16') - def _write_particles(self, ab: str, particles): + def _write_particles(self, index: int, ab: str, particles): """Write particles to HDF. data to write: x,y,z,size,source_intensity,max_image_photons,image_electrons,image_quantized_electrons,flag @@ -222,99 +246,117 @@ def _write_particles(self, ab: str, particles): n = len(particles) if particle_group not in self._h5: gr = self._h5.create_group(particle_group) - ds = gr.create_dataset('x', shape=(1, n), maxshape=(None, n), dtype='float32') + ds_shape = (self._n_images, n) + ds = gr.create_dataset('x', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.x - ds = gr.create_dataset('y', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('y', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.y - ds = gr.create_dataset('z', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('z', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.z - ds = gr.create_dataset('size', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('size', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.size - ds = gr.create_dataset('source_intensity', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('source_intensity', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.source_intensity - ds = gr.create_dataset('max_image_photons', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('max_image_photons', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.max_image_photons - ds = gr.create_dataset('image_electrons', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('image_electrons', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.image_electrons - ds = gr.create_dataset('image_quantized_electrons', shape=(1, n), maxshape=(None, n), dtype='float32') + ds = gr.create_dataset('image_quantized_electrons', shape=ds_shape, maxshape=(None, n), dtype='float32') ds[0, :] = particles.image_quantized_electrons - ds = gr.create_dataset('flag', shape=(1, n), maxshape=(None, n), dtype='uint8') + ds = gr.create_dataset('flag', shape=ds_shape, maxshape=(None, n), dtype='uint8') ds[0, :] = particles.flag ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - else: - ds = self._h5[particle_group]['x'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.x - ds = self._h5[particle_group]['y'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.y + return - ds = self._h5[particle_group]['z'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.z + ds = self._h5[particle_group]['x'] + curr_shape = ds.shape - ds = self._h5[particle_group]['size'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.size + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.x - ds = self._h5[particle_group]['source_intensity'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.source_intensity + ds = self._h5[particle_group]['y'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.y - ds = self._h5[particle_group]['max_image_photons'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.max_image_photons + ds = self._h5[particle_group]['z'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.z - ds = self._h5[particle_group]['image_electrons'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.image_electrons + ds = self._h5[particle_group]['size'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.size - ds = self._h5[particle_group]['image_quantized_electrons'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.image_quantized_electrons + ds = self._h5[particle_group]['source_intensity'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.source_intensity - ds = self._h5[particle_group]['flag'] - ds.resize((ds.shape[0] + 1, *ds.shape[1:])) - ds[-1, :] = particles.flag + ds = self._h5[particle_group]['max_image_photons'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.max_image_photons + + ds = self._h5[particle_group]['image_electrons'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.image_electrons + + ds = self._h5[particle_group]['image_quantized_electrons'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.image_quantized_electrons - def writeA(self, imgA: np.ndarray, particles: Particles = None): + ds = self._h5[particle_group]['flag'] + if self._n_images is None: + ds.resize((curr_shape[0] + 1, *curr_shape[1:])) + ds[index, :] = particles.flag + + def writeA(self, index: int, img: np.ndarray, particles: Particles = None): """Write image A Parameters ---------- - imgA: np.ndarray + index: int + Image index + img: np.ndarray Image A array particles: Particles Particle object used to generate img A """ ds = self._get_dsimgA() - ds.resize((self._image_index_a + 1, *ds.shape[1:])) - ds[self._image_index_a, ...] = imgA + if index >= ds.shape[0]: + raise KeyError(f'Image index {index} is out of range. Only {ds.shape[0]} images are expected, thus ' + f'index should be in the range [0, {ds.shape[0] - 1}]') + ds[index, ...] = img if particles is not None: - self._write_particles('A', particles) - - self._image_index_a += 1 + self._write_particles(index, 'A', particles) - def writeB(self, imgB: np.ndarray, particles: Particles = None): + def writeB(self, index: int, img: np.ndarray, particles: Particles = None): """Write image B Parameters ---------- - imgB: np.ndarray + index: int + Image index + img: np.ndarray Image B array particles: Particles Particle object used to generate img B """ ds = self._get_dsimgB() - ds.resize((self._image_index_a + 1, *ds.shape[1:])) - ds[self._image_index_a, ...] = imgB + if index >= ds.shape[0]: + raise KeyError(f'Image index {index} is out of range. Only {ds.shape[0]} images are expected, thus ' + f'index should be in the range [0, {ds.shape[0] - 1}]') + ds[index, ...] = img if particles is not None: - self._write_particles('B', particles) - - self._image_index_a += 1 + self._write_particles(index, 'B', particles) def __exit__(self, exc_type, exc_val, exc_tb): ds = self._get_dsimgA() diff --git a/synpivimage/particles.py b/synpivimage/particles.py index 97f0ead..12d4690 100644 --- a/synpivimage/particles.py +++ b/synpivimage/particles.py @@ -284,48 +284,48 @@ def save_jsonld(self, filename: Union[str, pathlib.Path]): from .codemeta import get_package_meta filename = pathlib.Path(filename) # .with_suffix('.jsonld') - source_intensity = 0. if np.all(self.source_intensity == 0) else self.source_intensity.tolist() - max_image_photons = 0. if np.all(self.max_image_photons == 0) else self.max_image_photons.tolist() - image_electrons = 0. if np.all(self.image_electrons == 0) else self.image_electrons.tolist() + source_intensity = 0. if np.all(self.source_intensity == 0) else self.source_intensity + max_image_photons = 0. if np.all(self.max_image_photons == 0) else self.max_image_photons + image_electrons = 0. if np.all(self.image_electrons == 0) else self.image_electrons image_quantized_electrons = 0. if np.all( - self.image_quantized_electrons == 0) else self.image_quantized_electrons.tolist() + self.image_quantized_electrons == 0) else self.image_quantized_electrons hasParameter = [ m4i.variable.NumericalVariable( label='x', - hasNumericalValue=self.x.tolist() + hasNumericalValue=self.x.astype("float16").tolist() ), m4i.variable.NumericalVariable( label='y', - hasNumericalValue=self.y.tolist() + hasNumericalValue=self.y.astype("float16").tolist() ), m4i.variable.NumericalVariable( label='z', - hasNumericalValue=self.z.tolist() + hasNumericalValue=self.z.astype("float16").tolist() ), m4i.variable.NumericalVariable( label='size', - hasNumericalValue=self.size.tolist() + hasNumericalValue=self.size.astype("float16").tolist() ), m4i.variable.NumericalVariable( label='flag', - hasNumericalValue=self.flag.tolist() + hasNumericalValue=self.flag.astype("uint8").tolist() ), m4i.variable.NumericalVariable( label='source_intensity', - hasNumericalValue=source_intensity + hasNumericalValue=np.asarray(source_intensity, dtype="uint16").tolist() ), m4i.variable.NumericalVariable( label='max_image_photons', - hasNumericalValue=max_image_photons + hasNumericalValue=np.asarray(max_image_photons, dtype="uint16").tolist() ), m4i.variable.NumericalVariable( label='image_electrons', - hasNumericalValue=image_electrons + hasNumericalValue=np.asarray(image_electrons, dtype="uint16").tolist() ), m4i.variable.NumericalVariable( label='image_quantized_electrons', - hasNumericalValue=image_quantized_electrons + hasNumericalValue=np.asarray(image_quantized_electrons, dtype="uint16").tolist() ), ] particles = pivmeta.SyntheticParticle( diff --git a/tests/test_core.py b/tests/test_core.py index 0376264..5410dfd 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -134,7 +134,7 @@ def test_generate_certain_ppp(self): ) target_ppp = 0.01 partA = Particles.generate(ppp=target_ppp, - dx_max=[100, 100], + dx_max=[10, 100], dy_max=[-100, 100], dz_max=[1, 1], camera=cam, laser=laser) @@ -299,6 +299,8 @@ def test_take_image(self): pathlib.Path(pathlib.Path('img01a.tiff')).unlink(missing_ok=True) pathlib.Path(pathlib.Path('img01a.hdf')).unlink(missing_ok=True) pathlib.Path(pathlib.Path('img01a.json')).unlink(missing_ok=True) + + # # # class TestCore(unittest.TestCase): diff --git a/tests/test_io.py b/tests/test_io.py index b21f199..2934744 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -114,11 +114,14 @@ def test_write_hdf(self): particle_peak_count=1000) with synpivimage.HDF5Writer(__this_dir__ / 'data.hdf', + n_images=2, overwrite=True, camera=cam, laser=laser) as h5: - h5.writeA(imgA=img) - h5.writeA(imgA=img) + h5.writeA(0, img=img) + h5.writeA(1, img=img) + with self.assertRaises(KeyError): + h5.writeA(2, img=img) with h5py.File(__this_dir__ / 'data.hdf', 'r') as h5: self.assertEqual(h5['images/img_A'].shape, (2, 16, 16))