diff --git a/align.py b/align.py index 28af8bdace995910fe3a97b10cd24e7a3a8678ad..41e71eace33af63d89750cd3dbf659a1197ce0f4 100644 --- a/align.py +++ b/align.py @@ -8,6 +8,7 @@ import os import dask.array as da from functools import partial from zarr_tools import convert +from zarr.errors import ArrayNotFoundError import yaml import fire from multiprocessing import Pool @@ -17,6 +18,7 @@ from scipy.ndimage import laplace, rotate import json + def align_multichip( BF_TRITC_2D_path, out_path, @@ -41,12 +43,6 @@ def align_multichip( else: rotation_data_deg = 0 - tif_paths = [ - f"{os.path.dirname(BF_TRITC_2D_path)}/{c}{unit}_aligned.tif" - for c in concentrations - ] - print("tif_paths: ", tif_paths) - data = read_dask(BF_TRITC_2D_path) template16 = tf.imread(template_path) big_labels = tf.imread(labels_path) @@ -59,15 +55,16 @@ def align_multichip( fit_poisson=fit_poisson, nmax=nmax, rotation_data_deg=rotation_data_deg, + aligned_path=out_path ) try: p = Pool(data.shape[0]) - out = p.map(fun, zip(data, tif_paths, concentrations)) + out = p.map(fun, zip(data, range(len(concentrations)), concentrations)) except TypeError as e: print(f"Pool failed due to {e.args}") - out = list(map(fun, zip(data, tif_paths, concentrations))) + out = list(map(fun, zip(data, concentrations))) finally: p.close() @@ -113,7 +110,7 @@ def align_parallel(args, **kwargs): def align2D( stack_dask, - path_tif, + chip_index, ab, template16=None, big_labels=None, @@ -121,16 +118,17 @@ def align2D( fit_poisson=True, nmax=20, rotation_data_deg=0, + aligned_path="" ): print(ab, unit) try: - aligned = tf.imread(path_tif) - print(f"already aligned: {path_tif}:{aligned.shape}") + aligned = da.from_zarr(os.path.join(aligned_path, "0"))[chip_index].compute() + print(f"already aligned: {aligned_path}:{aligned.shape}") counts = count_cells(aligned, ab=ab) intensity_table = get_intensity_table(aligned[2], aligned[1]) counts.loc[:, "intensity"] = intensity_table.intensity return {"stack": aligned, "counts": counts} - except FileNotFoundError: + except ArrayNotFoundError: print("Processing...") data = stack_dask.compute() @@ -152,7 +150,7 @@ def align2D( lambda_fit_result = poisson.fit( counts.query(f"n_cells < {nmax}").n_cells, title=f"automatic {ab}{unit.replace('_','/')}", - save_fig_path=path_tif.replace(".tif", "-counts-hist.png"), + save_fig_path=aligned_path.replace(".zarr", f"{ab}-counts-hist.png"), ) counts.loc[:, "poisson fit"] = lambda_fit_result