utils.py

from concurrent.futures import ThreadPoolExecutor
import math
import sys
import logging

import numpy as np
import tifffile
from imageio.v3 import imread
from tqdm import tqdm

from PyImarisWriter import PyImarisWriter as PW

"""
Imaris valid types mapping, currently only support uint8, uint16, uint32 and float32, if the input image is in other data type, it will be converted to the nearest supported data type before writing to Imaris file.
"""
imaris_valid_types = {
    np.dtype('uint8'): 'uint8',
    np.dtype('uint16'): 'uint16',
    np.dtype('uint32'): 'uint32',
    np.dtype('float32'): 'float32',
}

class MyCallbackClass(PW.CallbackClass):
    def __init__(self):
        self.mUserDataProgress = 0

    def RecordProgress(self, progress, total_bytes_written):
        progress100 = int(progress * 100)
        if progress100 - self.mUserDataProgress >= 1: # log every 1% progress
            self.mUserDataProgress = progress100
            print(f'User Progress {self.mUserDataProgress}%, Bytes written: {total_bytes_written}')


"""
Image reader backend mapping, currently support tifffile, cpptiff and imageio, you can choose the backend by passing the backend name to the read_image function, if the backend is not supported, it will raise a ValueError.
"""
image_reader_backend_mapping = {
    'tifffile': tifffile.imread,
    'imageio': imread
}

def split_files_into_blocks(file_list, block_size):
    blocks = []
    for i in range(0, len(file_list), block_size):
        blocks.append(file_list[i:i + block_size])
    return blocks

def read_image(fp, backend='tifffile'):
    _image_reader = image_reader_backend_mapping.get(backend, None)
    if _image_reader is None:
        raise ValueError(f"Unsupported image reader backend: {backend}, supported backends are: {list(image_reader_backend_mapping.keys())}")
    return _image_reader(fp)

def read_image_stack(file_list, thread_count=8, backend='tifffile'):
    _image_reader = image_reader_backend_mapping.get(backend, None)
    if _image_reader is None:
        raise ValueError(f"Unsupported image reader backend: {backend}, supported backends are: {list(image_reader_backend_mapping.keys())}")

    temp = _image_reader(file_list[0])
    if temp.ndim != 2:
        raise TypeError("Currently only support granyscale images.")
    Z = len(file_list)
    Y, X = temp.shape
    dtype = temp.dtype
    image_stack = np.empty((Z, Y, X), dtype=dtype)

    pbar = tqdm(total=len(file_list), desc="Loading images", unit="img", file=sys.stdout)
    def _imread(idx):
        image_stack[idx] = _image_reader(file_list[idx])
        pbar.update(1)

    with ThreadPoolExecutor(max_workers=thread_count) as executor:
        for i in range(len(file_list)):
            executor.submit(_imread, i)

    image_stack = np.asarray(image_stack)
    pbar.close()
    return image_stack

def block_divide(image_size, block_size):
    return PW.ImageSize(
        x=math.ceil(image_size.x / block_size.x),
        y=math.ceil(image_size.y / block_size.y),
        z=math.ceil(image_size.z / block_size.z),
        c=math.ceil(image_size.c / block_size.c),
        t=math.ceil(image_size.t / block_size.t)
    )

"""
Logging system
"""
def setup_logging(log_file):

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(log_file, mode='w'),
            logging.StreamHandler(sys.stdout)
        ]
    )

    class LoggerWriter:
        def __init__(self, level):
            self.level = level

        def write(self, msg):
            msg = msg.strip()
            if msg:
                logging.log(self.level, msg)

        def flush(self):
            pass

    sys.stdout = LoggerWriter(logging.INFO)
    sys.stderr = LoggerWriter(logging.ERROR)

    sys.excepthook = lambda t, v, tb: logging.error("Uncaught exception", exc_info=(t, v, tb))

"""
Timing function
"""
def format_time(seconds):
    if seconds < 60:
        return f"{seconds:.2f} s"
    elif seconds < 3600:
        m, s = divmod(seconds, 60)
        return f"{int(m)} m {s:.1f} s"
    else:
        h, rem = divmod(seconds, 3600)
        m, s = divmod(rem, 60)
        return f"{int(h)} h {int(m)} m {s:.1f} s"


if __name__ == "__main__":
    import glob
    from time import time

    file_pattern = r'test_image\z*.tif'
    file_list = sorted(glob.glob(file_pattern))

    start_time = time()
    image_stack = read_image_stack(file_list, thread_count=8)
    end_time = time()
    print(f"Time taken to read image stack: {end_time - start_time:.2f} seconds")