fix issues

examples/mousebrain_processing.py

@@ -35,6 +35,10 @@ def main(run_nnunet: bool = True, run_coiled_process: bool = True):
     if SUBJECT_ID == 'F4A1Te3Blaze':
         OME_ZARR_PATH = 'gcs://khanlab-lightsheet/data/mouse_appmaptapoe/bids/sub-F4A1Te3/micr/sub-F4A1Te3_sample-brain_acq-blaze4x_SPIM.ome.zarr'
 
+    RUN_ON_FULL_IM = False
+    if not RUN_ON_FULL_IM:
+        BA_CHANNEL = np.s_[0, 512:768, :, :]
+
     TARGET_PATH = f'C:/Users/than83/Documents/progtools/datasets/lightsheet_downsample/sub-{SUBJECT_ID}.ome.zarr'  # first downsample
     TARGET_FOLDER = f'C:/Users/than83/Documents/progtools/datasets/annotated/canvas_{SUBJECT_ID}'  # second and third downsample; other images
 
@@ -105,6 +109,8 @@ def main(run_nnunet: bool = True, run_coiled_process: bool = True):
         f'{TARGET_FOLDER}/.temp',
         GCS_BIAS_PATH
     )
+
+    bias_to_im_upscale = (4, 8, 8)
     async def fn(dask_worker):
         return await cvpl_nnunet_api.mousebrain_forward(
             dask_worker=dask_worker,
@@ -113,14 +119,12 @@ async def fn(dask_worker):
             NEG_MASK_PATH=GCS_NEG_MASK_TGT,
             GCS_BIAS_PATH=GCS_BIAS_PATH,
             BA_CHANNEL=BA_CHANNEL,
-            MAX_THRESHOLD=MAX_THRESHOLD
+            MAX_THRESHOLD=MAX_THRESHOLD,
+            bias_to_im_upscale=bias_to_im_upscale
         )
     cvpl_nnunet_api.coiled_run(fn=fn, nworkers=10, local_testing=False)
 
 
-
-
-
 if __name__ == '__main__':
     main()
 

examples/test_viewer.py

@@ -0,0 +1,17 @@
+import napari
+import cvpl_tools.tools.fs as tlfs
+import cvpl_tools.ome_zarr.napari.add as nozadd
+
+
+if __name__ == '__main__':
+    viewer = napari.Viewer(ndisplay=2)
+    nozadd.group_from_path(viewer, 'gcs://khanlab-lightsheet/data/mouse_appmaptapoe/bids/sub-F4A1Te3/micr/sub-F4A1Te3_sample-brain_acq-blaze4x_SPIM.ome.zarr', kwargs=dict(
+        name='im'
+    ))
+    # nozadd.group_from_path(viewer, 'gcs://khanlab-scratch/tmp/CacheDirectory_F4A1Te3Blaze/GLOBAL_LABEL/os/global_os/dask_im',
+    #                        kwargs=dict(
+    #                            name='os'
+    #                        ))
+    viewer.show(block=True)
+
+

src/cvpl_tools/im/algs/dask_label.py

@@ -158,7 +158,8 @@ def to_max(block: npt.NDArray, block_info: dict):
 
     locally_labeled = await tlfs.cache_im(
         lambda: im.map_blocks(map_block, meta=np.zeros(tuple(), dtype=output_dtype)),
-        context_args=dict(cache_url=None if fs is None else fs['locally_labeled_without_cumsum'], viewer_args=vargs)
+        context_args=context_args | dict(cache_url=None if fs is None else fs['locally_labeled_without_cumsum'],
+                                         viewer_args=vargs)
     )
 
     async def compute_nlbl_np_arr():
@@ -212,7 +213,8 @@ def compute_slices(block: npt.NDArray, block2: npt.NDArray, block_info: dict = N
 
     locally_labeled = await tlfs.cache_im(
         compute_locally_labeled,
-        context_args=dict(cache_url=None if fs is None else fs['locally_labeled_with_cumsum'], viewer_args=vargs)
+        context_args=context_args | dict(cache_url=None if fs is None else fs['locally_labeled_with_cumsum'],
+                                         viewer_args=vargs)
     )
 
     comp_i = 0
@@ -256,7 +258,7 @@ def local_to_global(block, block_info, ind_map_scatter):
 
     globally_labeled = await tlfs.cache_im(
         fn=compute_globally_labeled,
-        context_args=dict(
+        context_args=context_args | dict(
             cache_url=None if fs is None else fs['globally_labeled'],
             viewer_args=vargs
         )
@@ -269,7 +271,7 @@ def local_to_global(block, block_info, ind_map_scatter):
     if is_logging:
         print('Function ends')
 
-    im = await tlfs.cache_im(lambda: result_arr, context_args=dict(
+    im = await tlfs.cache_im(lambda: result_arr, context_args=context_args | dict(
         cache_url=None if fs is None else fs['global_os'],
         viewer_args=vargs | dict(is_label=True)
     ))

src/cvpl_tools/nnunet/api.py

@@ -72,7 +72,8 @@ async def mousebrain_forward(dask_worker,
                              NEG_MASK_PATH: str,
                              GCS_BIAS_PATH: str,
                              BA_CHANNEL: int,
-                             MAX_THRESHOLD: float
+                             MAX_THRESHOLD: float,
+                             bias_to_im_upscale: tuple,
                              ):
     # passing of dask_worker is credit to fjetter at https://github.com/dask/distributed/issues/8152
     from dask.distributed import Worker
@@ -308,7 +309,7 @@ async def fn(im, context_args):
     storage_options = dict(
         dimension_separator='/',
         preferred_chunksize=(2, 4096, 4096),
-        multiscale=3,
+        multiscale=4,
         compressor=numcodecs.Blosc(cname='lz4', clevel=9, shuffle=numcodecs.Blosc.BITSHUFFLE)
     )
     viewer_args = dict(
@@ -324,12 +325,12 @@ async def fn(im, context_args):
     async def compute_per_pixel_multiplier():
         with RDirFileSystem(NEG_MASK_PATH).open('', mode='rb') as infile:
             neg_mask = tifffile.imread(infile)
+        with RDirFileSystem(GCS_BIAS_PATH).open('', mode='rb') as infile:
+            bias = tifffile.imread(infile)
         neg_mask = dask_ndinterp.scale_nearest(da.from_array(neg_mask, chunks=(16, 16, 16)),
                                                scale=(2, 2, 2),
-                                               output_shape=neg_mask.shape,
+                                               output_shape=bias.shape,
                                                output_chunks=(32, 32, 32)).persist()
-        with RDirFileSystem(GCS_BIAS_PATH).open('', mode='rb') as infile:
-            bias = tifffile.imread(infile)
         bias = da.from_array(bias, chunks=(32, 32, 32))
         bias = dask_ndinterp.scale_nearest(bias,
                                            scale=(2, 2, 2), output_shape=neg_mask.shape, output_chunks=(64, 64, 64))
@@ -343,7 +344,7 @@ async def compute_per_pixel_multiplier():
                                )))
 
     async def compute_masking():
-        im = cur_im * dask_ndinterp.scale_nearest(ppm, scale=(4, 8, 8),
+        im = cur_im * dask_ndinterp.scale_nearest(ppm, scale=bias_to_im_upscale,
                                                   output_shape=cur_im.shape, output_chunks=(256, 512, 512))
         im = (im / MAX_THRESHOLD).clip(0., 1.)
         return im.astype(np.float16)

src/cvpl_tools/nnunet/current_im.py

@@ -1,6 +1,7 @@
 import copy
 
 import napari
+from mpmath.libmp import BACKEND
 
 import cvpl_tools.ome_zarr.io as ome_zarr_io
 import cvpl_tools.im.algs.dask_ndinterp as dask_ndinterp