parallel_read_to_imagej.py

# Open a tiff or h5 file in ImageJ by loading it to memory in several threads
######################################################################################
# Installation:

# 1. conda create -n pyimagej -c conda-forge pyimagej openjdk=8 scyjava dxchange h5py ipython
# 2. conda activate pyimagej
#######################################################################################
# Test:

# (pyimagej) usertxm@txmthree ~/vnikitin $ ipython -i parallel_read_to_imagej.py 

# In [1]: read_h5('/local/ssd/data/tmp_rec/usbat_sshc_pam_spot2_926_rec.h5')
# nproc=8, load time 15.498464107513428, load speed 2445.4036049692586MB/s

# In [2]: read_tiff('/local/ssd/data/tmp_rec/usbat_sshc_pam_spot2_926_rec/recon_00000.tiff')
# nproc=8, load time 13.721449136734009, load speed 2762.0989315579673MB/s

# In [3]: read_h5('/local/ssd/data/tmp_rec/usbat_sshc_pam_spot2_926_rec.h5',1)
# nproc=1, load time 42.046916246414185, load speed 901.3740693345653MB/s

# In [4]: read_tiff('/local/ssd/data/tmp_rec/usbat_sshc_pam_spot2_926_rec/recon_00000.tiff',1)
# nproc=1, load time 50.877564907073975, load speed 744.9255888960679MB/s
##################################################################

import imagej
import scyjava
import time
import dxchange
import h5py
import threading
import os
import numpy as np

scyjava.config.add_option('-Xmx512g')# limit memory
ij = imagej.init('/home/beams/USERTXM/Software/Fiji.app',mode='interactive')
ij.ui().showUI() # if you want to display the GUI immediately
from jpype import JClass
WindowManager = JClass('ij.WindowManager')


def read_h5_chunk(rec, fid, k, lchunk):
    """Read a chunk of data from h5"""
    st = k*lchunk
    end = min((k+1)*lchunk,rec.shape[0])
    rec[st:end] = fid[st:end]  

def read_tiff_chunk(rec, fname, k, lchunk):
    """Read a chunk of data from tiff"""
    st = k*lchunk
    end = min((k+1)*lchunk,rec.shape[0])
    rec[st:end] = dxchange.read_tiff_stack(fname,ind=range(st,end))

def read_h5(fname,nproc=8):
    t = time.time()
    with h5py.File(fname,'r') as fid:
        nz,n = fid['exchange/recon'].shape[:2]
        dtype = fid['exchange/recon'].dtype
        rec = np.zeros([nz,n,n], dtype=dtype)
        lchunk = int(np.ceil(nz/nproc))
        procs = []
        for k in range(nproc):
            read_proc = threading.Thread(
                target=read_h5_chunk, args=(rec, fid['exchange/recon'], k, lchunk))
            procs.append(read_proc)
            read_proc.start()
        for proc in procs:
            proc.join()
    tt = time.time()
    print(f'Load time {tt-t}')
    print(f'{nproc=}, load time {tt-t}, load speed {nz*n*n*4/(tt-t)/1024/1024}MB/s')


    ij.ui().show('recon', ij.py.to_dataset(rec))

def read_tiff(fname,nproc=8):
    t = time.time()
        
    tiff_file_list = sorted(
            list(filter(lambda x: x.endswith(('.tif', '.tiff')), os.listdir(os.path.dirname(fname)))))        
    z_start = int(tiff_file_list[0].split('.')[0].split('_')[1])        
    z_end = int(tiff_file_list[-1].split('.')[0].split('_')[1]) + 1        
    d = dxchange.read_tiff(fname)
    n = d.shape[-1]
    nz = z_end- z_start
    rec = np.zeros([nz,n,n], dtype=d.dtype)
    lchunk = int(np.ceil(nz/nproc))
    procs = []
    for k in range(nproc):
        read_proc = threading.Thread(
            target=read_tiff_chunk, args=(rec, fname, k, lchunk))
        procs.append(read_proc)
        read_proc.start()
    for proc in procs:
        proc.join()
    
    tt = time.time()
    print(f'{nproc=}, load time {tt-t}, load speed {nz*n*n*4/(tt-t)/1024/1024}MB/s')

    ij.ui().show('recon', ij.py.to_dataset(rec))