Multiscale ome-zarr v3 writer with sharding support

classes_mesoSPIM.mmd

@@ -79,6 +79,8 @@ classDiagram
   }
   class FilenameWizardSingleHDF5SelectionPage {
   }
+  FilenameWizardSingleOmeZarrSelectionPage {
+  }
   class FilenameWizardTiffSelectionPage {
   }
   class FilenameWizardWelcomePage {
@@ -264,6 +266,7 @@ classDiagram
   FilenameWizardBigTiffSelectionPage --|> AbstractSelectionPage
   FilenameWizardSingleHDF5SelectionPage --|> AbstractSelectionPage
   FilenameWizardTiffSelectionPage --|> AbstractSelectionPage
+  FilenameWizardSingleOmeZarrSelectionPage  --|> AbstractSelectionPage
   FifthChannelPage --|> GenericChannelPage
   FirstChannelPage --|> GenericChannelPage
   FourthChannelPage --|> GenericChannelPage

mesoSPIM/config/demo_config.py

@@ -338,6 +338,39 @@
         'transpose_xy': False, # in case X and Y axes need to be swapped for the correct tile positions
         }
 
+'''
+OME.ZARR parameters
+This will write a ome.zarr v3 multiscale on the fly during acquisition.
+The default parameter should work pretty well for most setups with little to no performance degradation
+during acquisition. Defaults include compression which will save disk space and can also improve
+performance because less data is written to disk. Data are written into shards which limits the number of
+files generated on disk. 
+
+Chunks can be set to adjust with each multiscale. Base and target chunks are defined and will start 
+with the base shape and automatically shift towards target with each scale. Chunks have a big influence on IO.
+Bigger chunks means less and more efficient IO, very small chunks will degrade performance on some hardware.
+
+compression: default: zstd-5. This is a good trade off of compute and compression. In our tests, there is 
+little to no performance degradation when using this setting.
+
+shards are defined by default. Be careful, shard shape must be defined carefully to prevent performance 
+degradation. We suggest that shards are shallow in Z and as large as you camera sensor in XY. 
+For best performance set the base and target chunks to the same z-depth as your shards.
+
+async_finalize, default True: Enables acquisition of the next tile to proceed immediately while the multiscale 
+is finalized in the background. On systems with slow IO, data can accumulate in RAM and cause a crash.
+Slow IO can be improved by using bigger chunks. If bigger chunks do not help, use async_finalize: False 
+to make mesoSPSIM pause after each tile acquisition until the multiscale is finished generating. 
+'''
+ome_zarr = {
+    'compression': 'zstd', # None, 'zstd', 'lz4'
+    'compression_level': 5, # 1-9
+    'shards': (64,6000,6000), # Specify Max shard size: suggest None for best performance, (64, 6000,6000) (axes: z,y,x)
+    'base_chunks': (64,256,256), # Starting chunk size (level 0). Bigger chunks, less files (axes: z,y,x)
+    'target_chunks': (64,64,64), # Approx ending chunks shape (level 5). Bigger chunks, less files (axes: z,y,x)
+    'async_finalize': True, # True, False
+    }
+
 '''
 Rescale the galvo amplitude when zoom is changed
 For example, if 'galvo_l_amplitude' = 1 V at zoom '1x', it will ve 2 V at zoom '0.5x'

mesoSPIM/scripts/run_cors_server.py

@@ -0,0 +1,21 @@
+#! python run_cors_server.py 8000 D:\MyFolder
+from http.server import SimpleHTTPRequestHandler, HTTPServer
+
+class CORSRequestHandler(SimpleHTTPRequestHandler):
+    def end_headers(self):
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
+        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
+        super().end_headers()
+
+if __name__ == '__main__':
+    import sys
+    import os
+
+    # Example: python cors_server.py 8000 D:\\MyFolder
+    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
+    directory = sys.argv[2] if len(sys.argv) > 2 else '.'
+    os.chdir(directory)
+    httpd = HTTPServer(('', port), CORSRequestHandler)
+    print(f"Serving {directory} at http://localhost:{port}")
+    httpd.serve_forever()

mesoSPIM/src/mesoSPIM_ImageWriter.py

@@ -3,6 +3,7 @@
 '''
 
 import os
+from pathlib import Path
 import time
 import numpy as np
 import tifffile
@@ -14,6 +15,11 @@
 import npy2bdv
 from .utils.acquisitions import AcquisitionList, Acquisition
 from .utils.utility_functions import write_line, gb_size_of_array_shape, replace_with_underscores, log_cpu_core
+from .utils.omezarr_writer import (PyramidSpec, ChunkScheme,
+                                   Live3DPyramidWriter, plan_levels,
+                                   compute_xy_only_levels, FlushPad,
+                                   BloscCodec, BloscShuffle
+                                   )
 
 
 class mesoSPIM_ImageWriter(QtCore.QObject):
@@ -39,7 +45,9 @@ def __init__(self, parent, frame_queue):
         self.y_pixels = int(self.y_pixels / self.y_binning)
 
         self.file_extension = ''
-        self.bdv_writer = self.tiff_writer = self.tiff_mip_writer = self.mip_image = None
+        self.bdv_writer = self.tiff_writer\
+            = self.tiff_mip_writer = self.mip_image\
+            = self.omezarr_writer = None
         self.tiff_aliases = ('.tif', '.tiff')
         self.bigtiff_aliases = ('.btf', '.tf2', '.tf8')
         self.check_versions()
@@ -64,9 +72,10 @@ def check_versions(self):
     def prepare_acquisition(self, acq, acq_list):
         self.folder = acq['folder']
         self.filename = replace_with_underscores(acq['filename'])
-        self.path = os.path.realpath(self.folder+'/'+self.filename)
+        self.path = os.path.realpath(self.folder+'/ '+ self.filename)
         self.MIP_path = os.path.realpath(self.folder +'/MAX_'+ self.filename + '.tiff')
-        self.file_root, self.file_extension = os.path.splitext(self.path)
+        # self.file_root, self.file_extension = os.path.splitext(self.path)
+        self.file_extension = ''.join(Path(self.path).suffixes)
         logger.info(f'Save path: {self.path}')
 
         self.binning_string = self.state['camera_binning'] # Should return a string in the form '2x4'
@@ -77,7 +86,63 @@ def prepare_acquisition(self, acq, acq_list):
         self.y_pixels = int(self.y_pixels / self.y_binning)
         self.max_frame = acq.get_image_count()
 
-        if self.file_extension == '.h5':
+        if acq == acq_list[0]:
+            self.first_path = self.path
+
+        print(f'{self.file_extension=}')
+        if self.file_extension == '.ome.zarr':
+            if acq == acq_list[0]:
+                import zarr
+                zarr.open_group(self.first_path, mode="a")
+            if hasattr(self.cfg, "ome_zarr"):
+                compression = self.cfg.ome_zarr['compression']
+                compression_level = self.cfg.ome_zarr['compression_level']
+                shards = self.cfg.ome_zarr['shards']
+                base_chunks = self.cfg.ome_zarr['base_chunks']
+                target_chunks = self.cfg.ome_zarr['target_chunks']
+                async_finalize = self.cfg.ome_zarr['async_finalize']
+            else:
+                compression = 'zstd'
+                compression_level = 5
+                shards = None
+                base_chunks = (256,256,256)
+                target_chunks = (256,256,256)
+                async_finalize = True
+
+            Z_EST, Y, X = (self.max_frame, self.x_pixels, self.y_pixels)
+            px_size_zyx = (acq['z_step'], self.cfg.pixelsize[acq['zoom']], self.cfg.pixelsize[acq['zoom']])
+
+            xy_levels = compute_xy_only_levels(px_size_zyx)
+            levels = plan_levels(Y, X, Z_EST, xy_levels, min_dim=64)
+
+            spec = PyramidSpec(
+                z_size_estimate=Z_EST,  # big upper bound; we'll truncate at the end
+                y=Y, x=X, levels=levels,
+            )
+
+            shard_shape = shards
+            scheme = ChunkScheme(base=base_chunks, target=target_chunks)
+
+            compressor = compression
+            if compression:
+                compressor = BloscCodec(cname=compression, clevel=compression_level, shuffle=BloscShuffle.bitshuffle)
+
+            isetup = acq_list.index(acq)
+            group_name = 's{:d}-t{:d}.zarr'.format(isetup, 0) # time = 0 for now
+            self.omezarr_writer = Live3DPyramidWriter(
+                    spec,
+                    voxel_size=px_size_zyx,
+                    path=self.first_path + '/' + group_name,
+                    max_workers=os.cpu_count() // 2,
+                    chunk_scheme=scheme,
+                    compressor=compressor,
+                    shard_shape=shard_shape,
+                    flush_pad=FlushPad.DUPLICATE_LAST,  # keeps alignment, no RMW
+                    async_close=async_finalize,
+                    translation=(acq['z_start'], acq['y_pos'], acq['x_pos'])
+            )
+
+        elif self.file_extension == '.h5':
             if hasattr(self.cfg, "hdf5"):
                 subsamp = self.cfg.hdf5['subsamp']
                 compression = self.cfg.hdf5['compression']
@@ -161,7 +226,10 @@ def image_to_disk(self, acq, acq_list, image):
         if self.cur_image_counter % 5 == 0:
             self.parent.sig_status_message.emit('Writing to disk...')
         xy_res = (1./self.cfg.pixelsize[acq['zoom']], 1./self.cfg.pixelsize[acq['zoom']])
-        if self.file_extension == '.h5':
+
+        if self.file_extension == '.ome.zarr':
+            self.omezarr_writer.push_slice(image)
+        elif self.file_extension == '.h5':
             self.bdv_writer.append_plane(plane=image, z=self.cur_image_counter,
                                             illumination=acq_list.find_value_index(acq['shutterconfig'], 'shutterconfig'),
                                             channel=acq_list.find_value_index(acq['laser'], 'laser'),
@@ -193,7 +261,9 @@ def abort_writing(self):
         self.abort_flag = True
         if self.running_flag:
             try:
-                if self.file_extension == '.h5':
+                if self.file_extension == '.ome.zarr':
+                    self.omezarr_writer.__exit__(None, None, None)
+                elif self.file_extension == '.h5':
                     self.bdv_writer.close()
                 elif self.file_extension == '.raw':
                     del self.xy_stack
@@ -211,7 +281,9 @@ def abort_writing(self):
     @QtCore.pyqtSlot(Acquisition, AcquisitionList)
     def end_acquisition(self, acq, acq_list):
         logger.info("end_acquisition() started")
-        if self.file_extension == '.h5':
+        if self.file_extension == '.ome.zarr':
+            self.omezarr_writer.__exit__(None, None, None)
+        elif self.file_extension == '.h5':
             if acq == acq_list[-1]:
                 try:
                     self.bdv_writer.set_attribute_labels('channel', tuple(acq_list.get_unique_attr_list('laser')))
@@ -303,7 +375,7 @@ def write_snap_metadata(self, path):
     def write_metadata(self, acq, acq_list):
         logger.debug("write_metadata() started")
         ''' Writes a metadata.txt file. Path contains the file to be written '''
-        path = acq['folder'] + '/' + acq['filename']
+        path = acq['folder'] + '/' + self.filename
         metadata_path = os.path.dirname(path) + '/' + os.path.basename(path) + '_meta.txt'
 
         if acq['filename'][-3:] == '.h5':

mesoSPIM/src/utils/acquisitions.py

@@ -4,7 +4,7 @@
 
 Helper classes for mesoSPIM acquisitions
 '''
-
+from pathlib import Path
 import indexed
 import os.path
 import logging
@@ -336,7 +336,8 @@ def check_for_duplicated_filenames(self):
         filenames = []
         # Create a list of full file paths
         for i in range(len(self)):
-            if self[i]['filename'][-3:] != '.h5':
+            file_extension = ''.join(Path(self[i]['filename']).suffixes)
+            if file_extension not in ('.h5', '.ome.zarr'):
                 filename = self[i]['folder']+'/'+self[i]['filename']
                 filenames.append(filename)
         duplicates = self.get_duplicates_in_list(filenames)

mesoSPIM/src/utils/filename_wizard.py

@@ -14,8 +14,8 @@
 class FilenameWizard(QtWidgets.QWizard):
     wizard_done = QtCore.pyqtSignal()
 
-    num_of_pages = 6
-    (welcome, raw, tiff, bigtiff, single_hdf5, finished) = range(num_of_pages)
+    num_of_pages = 7
+    (welcome, raw, tiff, bigtiff, single_hdf5, omezarr_string, finished) = range(num_of_pages)
 
     def __init__(self, parent=None):
         '''Parent is object of class mesoSPIM_AcquisitionManagerWindow()'''
@@ -25,14 +25,15 @@ def __init__(self, parent=None):
         through '''
         self.parent = parent
         self.state = self.parent.state # the mesoSPIM_StateSingleton() instance
-        self.file_format = None  # 'raw', 'h5', 'tiff', 'btf'
+        self.file_format = None  # 'raw', 'h5', 'tiff', 'btf', 'ome.zarr'
         self.setWindowTitle('Filename Wizard')
         self.setPage(0, FilenameWizardWelcomePage(self))
         self.setPage(1, FilenameWizardRawSelectionPage(self))
         self.setPage(2, FilenameWizardTiffSelectionPage(self))
         self.setPage(3, FilenameWizardBigTiffSelectionPage(self))
         self.setPage(4, FilenameWizardSingleHDF5SelectionPage(self))
-        self.setPage(5, FilenameWizardCheckResultsPage(self))
+        self.setPage(5, FilenameWizardSingleOmeZarrSelectionPage(self))
+        self.setPage(6, FilenameWizardCheckResultsPage(self))
         self.setStyleSheet(''' font-size: 16px; ''')
         self.show()
 
@@ -116,14 +117,26 @@ def generate_filename_list(self, increment_number=True):
 
                 file_suffix = '.' + self.file_format
 
-            elif self.file_format == 'h5':
+            elif self.file_format in ('h5', 'ome.zarr'):
                 if self.field('DescriptionHDF5'):
                     filename += replace_with_underscores(self.field('DescriptionHDF5')) + '_'
+                    file_suffix = '_bdv.' + self.file_format
+                elif self.field('DescriptionOmeZarr'):
+                    filename += replace_with_underscores(self.field('DescriptionOmeZarr')) + '_'
+                    file_suffix = '.' + self.file_format
                 filename += f'Mag{self.parent.model.getZoom(0)}'
                 laser_list = self.parent.model.getLaserList()
                 for laser in laser_list:
                     filename += '_ch' + laser[:-3]
-                file_suffix = '_bdv.' + self.file_format
+
+            # elif self.file_format == 'ome.zarr':
+            #     if self.field('DescriptionOmeZarr'):
+            #         filename += replace_with_underscores(self.field('DescriptionOmeZarr')) + '_'
+            #     filename += f'Mag{self.parent.model.getZoom(0)}'
+            #     laser_list = self.parent.model.getLaserList()
+            #     for laser in laser_list:
+            #         filename += '_ch' + laser[:-3]
+            #     file_suffix = '.' + self.file_format
 
             else:
                 raise ValueError(f"file suffix invalid: {self.file_format}")
@@ -153,11 +166,13 @@ def __init__(self, parent=None):
         self.tiff_string = 'ImageJ TIFF files: ~.tiff'
         self.bigtiff_string = 'BigTIFF files: ~.btf'
         self.single_hdf5_string = 'BigDataViewer HDF5 file: ~.h5'
+        self.omezarr_string = 'OME-ZARR: ~.ome.zarr'
 
         self.SaveAsComboBoxLabel = QtWidgets.QLabel('Save as:')
         self.SaveAsComboBox = QtWidgets.QComboBox()
-        self.SaveAsComboBox.addItems([self.raw_string, self.tiff_string, self.bigtiff_string, self.single_hdf5_string])
-        self.SaveAsComboBox.setCurrentIndex(3)
+        self.SaveAsComboBox.addItems([self.raw_string, self.tiff_string,
+                                      self.bigtiff_string, self.single_hdf5_string, self.omezarr_string])
+        self.SaveAsComboBox.setCurrentIndex(4)
 
         self.registerField('SaveAs', self.SaveAsComboBox, 'currentIndex')
 
@@ -179,6 +194,9 @@ def nextId(self):
         elif self.SaveAsComboBox.currentText() == self.single_hdf5_string:
             self.parent.file_format = 'h5'
             return self.parent.single_hdf5
+        elif self.SaveAsComboBox.currentText() == self.omezarr_string:
+            self.parent.file_format = 'ome.zarr'
+            return self.parent.omezarr_string
 
 
 class AbstractSelectionPage(QtWidgets.QWizardPage):
@@ -285,6 +303,18 @@ def validatePage(self):
         self.parent.generate_filename_list(increment_number=False)
         return super().validatePage()
 
+class FilenameWizardSingleOmeZarrSelectionPage(AbstractSelectionPage):
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.setTitle("Autogenerate ome-zarr filename")
+        self.setSubTitle("All raw data saved into one ome-zarr file, accompanied by two metadata files."
+                         "\nFilename example: {Description}_Mag1x_ch488_ch561.ome.zarr")
+        self.registerField('DescriptionOmeZarr', self.DescriptionLineEdit)
+
+    # def validatePage(self):
+    #     self.parent.generate_filename_list(increment_number=False)
+    #     return super().validatePage()
+
 
 class FilenameWizardCheckResultsPage(QtWidgets.QWizardPage):
     def __init__(self, parent=None):
@@ -305,12 +335,13 @@ def __init__(self, parent=None):
         self.setLayout(self.layout)
 
     def initializePage(self):
-        if self.parent.file_format in ('raw', 'tiff', 'btf'):
+        if self.parent.file_format in ('raw', 'tiff', 'btf', 'ome.zarr'):
             file_list = self.parent.filename_list
         elif self.parent.file_format == 'h5':
             file_list = [self.parent.filename_list[0]]
         else:
-            raise ValueError(f"file_format must be in ('raw', 'tiff', 'btf', 'h5'), received {self.parent.file_format}")
+            print(self.parent.filename_list)
+            raise ValueError(f"file_format must be in ('raw', 'tiff', 'btf', 'h5' 'ome.zarr'), received {self.parent.file_format}")
 
         for f in file_list:
             self.mystring += f