Proof-of-concept lazy reader for Pyramidal-OME-TIFF as Zarr

ome_tiff_reader.py

import tifffile
import zarr
import numpy as np

class PyramidalOMETIFFStore:
    
    def __init__(self, path):
        self.tf = tifffile.TiffFile(path)
        self.fh = self.tf.filehandle
        
        if not self.tf.is_ome or not self.tf.series[0].is_pyramid:
            raise ValueError('File not a pyramidal OME-TIFF')
        
        self.pyramid = self.tf.series[0].levels
        
    def __getitem__(self, key):
        if key == '.zgroup':
            return { 
                "zarr_format": 2 
            }
        
        level_key, rest = key.split('/')
        level = self.pyramid[int(level_key)]
        if rest == '.zarray':
            first_page = level[0].aspage()
            dtype = self.tf.byteorder + first_page.dtype.str[1:]
            return {
                "chunks": [
                    first_page.tiledepth,
                    first_page.tilelength,
                    first_page.tilewidth
                ],
                "compressor": None,
                "dtype": dtype,
                "fill_value": 0.0,
                "filters": None,
                "order": "C",
                "shape": level.shape,
                "zarr_format": 2
            }
        
        c, y, x = [int(k) for k in rest.split('.')]
        page = level[c].aspage()
        arr = self._read_tile(page, x, y)
        return arr.tobytes()    
        
    def __setitem__(self, key, value):
        pass
    
    def keys(self):
        yield '.zgroup'
        for i in range(len(self.pyramid)):
            yield f'{i}/.zarray'
    
    def __iter__(self):
        return self.keys()
    
    def _read_tile(self, page, x, y):
        offsets = page.dataoffsets
        bytecounts = page.databytecounts
        xstride, ystride = np.ceil(
            np.array(page.shape) / page.tilewidth
        ).astype(int)
        index = x + y * ystride
        segment, = list(
            self.fh.read_segments(
                offsets[index:index+1], 
                bytecounts[index:index+1]
            )
        )
        data, index, shape = page.decode(*segment)
        return data.squeeze()