check for zarr output instead of tifs

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