diff --git a/MANIFEST.in b/MANIFEST.in index bc04769..b57e5a1 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,7 @@ include LICENSE include README.md exclude .pre-commit-config.yaml +exclude *.ipynb recursive-include mesospim_stitcher *.py recursive-include mesospim_stitcher *.ijm diff --git a/mesospim_stitcher/fuse.py b/mesospim_stitcher/fuse.py index f9dd1ea..8f53592 100644 --- a/mesospim_stitcher/fuse.py +++ b/mesospim_stitcher/fuse.py @@ -5,7 +5,9 @@ import dask.array as da import h5py +import numpy as np import zarr +from mpi4py import MPI from numcodecs import Blosc from ome_zarr.io import parse_url from ome_zarr.writer import write_image @@ -28,6 +30,12 @@ def fuse_image( grid_transforms = root.findall( ".//ViewTransform/[Name='Translation from Tile Configuration']/affine" ) + + if len(grid_transforms) == 0: + grid_transforms = root.findall( + ".//ViewTransform/[Name='Translation to Regular Grid']/affine" + ) + assert grid_transforms is not None, "No grid transforms found in XML file" z_scale_str = root.find(".//ViewTransform/[Name='calibration']/affine") @@ -104,6 +112,16 @@ def fuse_image( x_s, x_e, y_s, y_e, z_s, z_e = translations[i] new_image[z_s:z_e, y_s:y_e, x_s:x_e] = curr_tile + try: + # write_ome_zarr(output_path, new_image, overwrite) + write_hdf5(output_path, new_image, overwrite) + except Exception as e: + raise e + finally: + input_file.close() + + +def write_ome_zarr(output_path: Path, image: da, overwrite: bool): if output_path.exists() and overwrite: rmtree(output_path) else: @@ -116,7 +134,7 @@ def fuse_image( root = zarr.group(store=store) compressor = Blosc(cname="zstd", clevel=4, shuffle=Blosc.SHUFFLE) write_image( - image=new_image, + image=image, group=root, axes=[ {"name": "z", "type": "space", "unit": "micrometer"}, @@ -131,7 +149,7 @@ def fuse_image( [{"type": "scale", "scale": [10.0, 41.6, 41.6]}], ], storage_options=dict( - chunks=(2, new_image.shape[1], new_image.shape[2]), + chunks=(2, image.shape[1], image.shape[2]), compressor=compressor, ), ) @@ -146,4 +164,92 @@ def fuse_image( ] } - input_file.close() + +def write_hdf5(output_path: Path, image: da, overwrite: bool): + subdivisions = np.array( + [[32, 32, 16], [32, 32, 16], [32, 32, 16], [32, 32, 16], [32, 32, 16]] + ) + resolutions = np.array( + [[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8], [16, 16, 16]] + ) + + rank = MPI.COMM_WORLD.rank + # + # print(f"Rank {rank} starting to write") + f = h5py.File(output_path, "w") + + f.create_dataset( + "s00/subdivisions", + data=subdivisions, + shape=subdivisions.shape, + dtype="uint16", + ) + f.create_dataset( + "s00/resolutions", + data=resolutions, + shape=resolutions.shape, + dtype="uint16", + ) + + data_group: h5py.Group = f.create_group("t00000/s00") + + orig_image = data_group.create_dataset( + "0/cells", + data=image, + chunks=(16, 32, 32), + dtype="uint16", + shape=image.shape, + ) + + image_shape = image.shape + x_split_size = image_shape[2] // 4 + y_split_size = image_shape[1] // 3 + + x_borders = [ + 0, + x_split_size, + x_split_size * 2, + x_split_size * 3, + image_shape[2], + ] + y_borders = [0, y_split_size, y_split_size * 2, image_shape[1]] + + x_start = x_borders[rank % 4] + x_end = x_borders[rank % 4 + 1] + y_start = y_borders[rank // 4] + y_end = y_borders[(rank // 4 + 1)] + + orig_image[:, y_start:y_end, x_start:x_end] = image[ + :, y_start:y_end, x_start:x_end + ] + + for i in range(1, resolutions.shape[0]): + prev_resolution = data_group[f"{i-1}/cells"][ + :, y_start:y_end, x_start:x_end + ] + data_group.require_dataset( + f"{i}/cells", + data=prev_resolution[::2, ::2, ::2], + chunks=(16, 32, 32), + compression="gzip", + dtype="uint16", + shape=prev_resolution[::2, ::2, ::2].shape, + ) + + f.close() + + # print(f"Rank {rank} finished writing") + + +if __name__ == "__main__": + xml_path = Path( + "/mnt/Data/TiledDataset/2.5x_tile/" + "One_Channel/2.5x_tile_igor_rightonly_Mag2.5x_" + "ch488_ch561_ch647_bdv.xml" + ) + input_path = Path("/mnt/Data/TiledDataset/2.5x_tile/One_Channel/test.h5") + output_path = Path( + "/mnt/Data/TiledDataset/2.5x_tile/One_Channel/test_out.h5" + ) + + fuse_image(xml_path, input_path, output_path) diff --git a/playing_around.ipynb b/playing_around.ipynb new file mode 100644 index 0000000..29cd743 --- /dev/null +++ b/playing_around.ipynb @@ -0,0 +1,35 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "initial_id", + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}