Fix preprocessing with zarr-python 3 (#294)

* use tensorstore for grid sampling

* remove old image sampling

* fix channel enumeration

Author: ziw-liu

viscy/utils/meta_utils.py

@@ -4,9 +4,10 @@
 import iohub.ngff as ngff
 import numpy as np
 import pandas as pd
+import tensorstore
+from tqdm import tqdm
 
-import viscy.utils.mp_utils as mp_utils
-from viscy.utils.cli_utils import show_progress_bar
+from viscy.utils.mp_utils import get_val_stats
 
 
 def write_meta_field(position: ngff.Position, metadata, field_name, subfield_name):
@@ -45,11 +46,23 @@ def write_meta_field(position: ngff.Position, metadata, field_name, subfield_nam
         position.zattrs[field_name] = field_metadata
 
 
+def _grid_sample(
+    position: ngff.Position, grid_spacing: int, channel_index: int, num_workers: int
+):
+    return (
+        position["0"]
+        .tensorstore(
+            context=tensorstore.Context(
+                {"data_copy_concurrency": {"limit": num_workers}}
+            )
+        )[:, channel_index, :, ::grid_spacing, ::grid_spacing]
+        .read()
+        .result()
+    )
+
+
 def generate_normalization_metadata(
-    zarr_dir,
-    num_workers=4,
-    channel_ids=-1,
-    grid_spacing=32,
+    zarr_dir, num_workers=4, channel_ids=-1, grid_spacing=32
 ):
     """
     Generate pixel intensity metadata to be later used in on-the-fly normalization
@@ -89,54 +102,37 @@ def generate_normalization_metadata(
         mp_grid_sampler_args.append([position, grid_spacing])
 
     # sample values and use them to get normalization statistics
-    for i, channel in enumerate(channel_ids):
-        show_progress_bar(
-            dataloader=channel_ids,
-            current=i,
-            process="sampling channel values",
-        )
+    for i, channel_index in enumerate(channel_ids):
+        print(f"Sampling channel index {channel_index} ({i + 1}/{len(channel_ids)})")
 
-        channel_name = plate.channel_names[channel]
-        this_channels_args = tuple([args + [channel] for args in mp_grid_sampler_args])
+        channel_name = plate.channel_names[channel_index]
+        dataset_sample_values = []
+        position_and_statistics = []
 
-        # NOTE: Doing sequential mp with pool execution creates synchronization
-        #      points between each step. This could be detrimental to performance
-        positions, fov_sample_values = mp_utils.mp_sample_im_pixels(
-            this_channels_args, num_workers
-        )
-        dataset_sample_values = np.concatenate(
-            [arr.flatten() for arr in fov_sample_values]
-        )
-        fov_level_statistics = mp_utils.mp_get_val_stats(fov_sample_values, num_workers)
-        dataset_level_statistics = mp_utils.get_val_stats(dataset_sample_values)
+        for _, pos in tqdm(position_map, desc="Positions"):
+            samples = _grid_sample(pos, grid_spacing, channel_index, num_workers)
+            dataset_sample_values.append(samples)
+            fov_level_statistics = {"fov_statistics": get_val_stats(samples)}
+            position_and_statistics.append((pos, fov_level_statistics))
 
         dataset_statistics = {
-            "dataset_statistics": dataset_level_statistics,
+            "dataset_statistics": get_val_stats(np.stack(dataset_sample_values)),
         }
-
         write_meta_field(
             position=plate,
             metadata=dataset_statistics,
             field_name="normalization",
             subfield_name=channel_name,
         )
 
-        for j, pos in enumerate(positions):
-            show_progress_bar(
-                dataloader=position_map,
-                current=j,
-                process=f"calculating channel statistics {channel}/{list(channel_ids)}",
-            )
-            position_statistics = dataset_statistics | {
-                "fov_statistics": fov_level_statistics[j],
-            }
-
+        for pos, position_statistics in position_and_statistics:
             write_meta_field(
                 position=pos,
-                metadata=position_statistics,
+                metadata=dataset_statistics | position_statistics,
                 field_name="normalization",
                 subfield_name=channel_name,
             )
+
     plate.close()
 
 

viscy/utils/mp_utils.py

@@ -268,55 +268,3 @@ def get_val_stats(sample_values):
         "iqr": float(scipy.stats.iqr(sample_values)),
     }
     return meta_row
-
-
-def mp_sample_im_pixels(fn_args, workers):
-    """Read and computes statistics of images with multiprocessing
-
-    :param list of tuple fn_args: list with tuples of function arguments
-    :param int workers: max number of workers
-    :return: list of paths and corresponding returned df from get_im_stats
-    """
-
-    with ProcessPoolExecutor(workers) as ex:
-        # can't use map directly as it works only with single arg functions
-        res = ex.map(sample_im_pixels, *zip(*fn_args))
-    return list(map(list, zip(*list(res))))
-
-
-def sample_im_pixels(
-    position: ngff.Position,
-    grid_spacing,
-    channel,
-):
-    # TODO move out of mp utils into normalization utils
-    """
-    Read and computes statistics of images for each point in a grid.
-    Grid spacing determines distance in pixels between grid points
-    for rows and cols.
-    By default, samples from every time position and every z-depth, and
-    assumes that the data in the zarr store is stored in [T,C,Z,Y,X] format,
-    for time, channel, z, y, x.
-
-    :param Position zarr_dir: NGFF position node object
-    :param int grid_spacing: spacing of sampling grid in x and y
-    :param int channel: channel to sample from
-
-    :return list meta_rows: Dicts with intensity data for each grid point
-    """
-    image_zarr = position.data
-
-    all_sample_values = []
-    all_time_indices = list(range(image_zarr.shape[0]))
-    all_z_indices = list(range(image_zarr.shape[2]))
-
-    for time_index in all_time_indices:
-        for z_index in all_z_indices:
-            image_slice = image_zarr[time_index, channel, z_index, :, :]
-            _, _, sample_values = image_utils.grid_sample_pixel_values(
-                image_slice, grid_spacing
-            )
-            all_sample_values.append(sample_values)
-    sample_values = np.stack(all_sample_values, 0).flatten()
-
-    return position, sample_values