Merge pull request #6 from AllenNeuralDynamics/feat-pypi-release

Feat pypi release

Author: camilolaiton

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+exclude =
+    .git,
+    __pycache__,
+    build
+max-complexity = 10
+max-line-length = 100

README.md

@@ -1,6 +1,6 @@
-# aind-z1-cell-segmentation
+# aind-large-scale-cellpose
 
-Large-scale cell segmentation using cellpose for Z1 data.
+Large-scale cell segmentation using cellpose.
 
 The approach is the following:
 

code/cell_segmentation/__init__.py code/aind_large_scale_cellpose/__init__.py

@@ -81,9 +81,9 @@ def execute_worker(
     )
 
     # Cell probability above threshold
-    cell_probability = (
-        data[0][-1] + data[1][-1] + data[2][-1] > cellprob_threshold
-    ).astype(np.uint8)
+    cell_probability = (data[0][-1] + data[1][-1] + data[2][-1] > cellprob_threshold).astype(
+        np.uint8
+    )
 
     # Looking at flows within cell areas
     dP_masked = dP * cell_probability
@@ -218,9 +218,7 @@ def combine_gradients(
         0,
         0,
     )
-    logger.info(
-        f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}"
-    )
+    logger.info(f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}")
 
     lazy_data = (
         ImageReaderFactory()
@@ -270,18 +268,14 @@ def combine_gradients(
     logger.info(
         f"Combined gradients: {output_combined_gradients} - chunks: {output_combined_gradients.chunks}"  # noqa: E501
     )
-    logger.info(
-        f"Cell probabilities path: {output_cellprob} - chunks: {output_cellprob.chunks}"
-    )
+    logger.info(f"Cell probabilities path: {output_cellprob} - chunks: {output_cellprob.chunks}")
 
     # Estimating total batches
     total_batches = np.prod(zarr_dataset.lazy_data.shape) / (
         np.prod(zarr_dataset.prediction_chunksize) * batch_size
     )
     samples_per_iter = n_workers * batch_size
-    logger.info(
-        f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}"
-    )
+    logger.info(f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}")
 
     logger.info(f"{20*'='} Starting combination of gradients {20*'='}")
     start_time = time()

code/cell_segmentation/cellpose_segmentation/__init__.py code/aind_large_scale_cellpose/cellpose_segmentation/__init__.py

@@ -15,8 +15,7 @@
 import zarr
 from aind_large_scale_prediction._shared.types import ArrayLike, PathLike
 from aind_large_scale_prediction.generator.dataset import create_data_loader
-from aind_large_scale_prediction.generator.utils import (
-    recover_global_position, unpad_global_coords)
+from aind_large_scale_prediction.generator.utils import recover_global_position, unpad_global_coords
 from aind_large_scale_prediction.io import ImageReaderFactory
 from cellpose import core
 from cellpose.dynamics import follow_flows
@@ -69,9 +68,7 @@ def computing_overlapping_hist_and_seed_finding(
     # Flatten p and compute edges
     p_flatten = p.astype("int32").reshape(p.shape[0], -1)
     shape0 = p.shape[1:]
-    edges = [
-        np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1) for i in range(len(shape0))
-    ]
+    edges = [np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1) for i in range(len(shape0))]
 
     # Compute histogram
     h, _ = np.histogramdd(tuple(p_flatten), bins=edges)
@@ -94,9 +91,7 @@ def computing_overlapping_hist_and_seed_finding(
 
     # Compute pixel coordinates
     pix_local = np.column_stack(seeds_sorted).astype(np.uint32)
-    pix_global = pix_local + np.array(
-        [global_coord.start for global_coord in global_coords]
-    )
+    pix_global = pix_local + np.array([global_coord.start for global_coord in global_coords])
 
     return pix_global, pix_local, h
 
@@ -336,12 +331,8 @@ def generate_flows_and_centroids(
         multiprocessing.set_start_method("spawn", force=True)
 
     # Getting overlap prediction chunksize
-    overlap_prediction_chunksize = (0,) + tuple(
-        [axis_overlap * 2] * len(prediction_chunksize[-3:])
-    )
-    logger.info(
-        f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}"
-    )
+    overlap_prediction_chunksize = (0,) + tuple([axis_overlap * 2] * len(prediction_chunksize[-3:]))
+    logger.info(f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}")
 
     lazy_data = (
         ImageReaderFactory()
@@ -406,9 +397,7 @@ def generate_flows_and_centroids(
         np.prod(zarr_dataset.prediction_chunksize) * batch_size
     )
     samples_per_iter = n_workers * batch_size
-    logger.info(
-        f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}"
-    )
+    logger.info(f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}")
 
     logger.info(f"{20*'='} Combining flows and creating histograms {20*'='}")
     start_time = time()
@@ -465,9 +454,7 @@ def generate_flows_and_centroids(
             curr_picked_blocks = 0
             picked_blocks = []
             time_proc_blocks_end = time()
-            logger.info(
-                f"Time processing blocks: {time_proc_blocks_end - time_proc_blocks}"
-            )
+            logger.info(f"Time processing blocks: {time_proc_blocks_end - time_proc_blocks}")
 
     if curr_picked_blocks != 0:
         logger.info(f"Blocks not processed inside of loop: {curr_picked_blocks}")

code/cell_segmentation/cellpose_segmentation/_shared/__init__.py code/aind_large_scale_cellpose/cellpose_segmentation/_shared/__init__.py

@@ -18,8 +18,7 @@
 import psutil
 import zarr
 from aind_large_scale_prediction.generator.dataset import create_data_loader
-from aind_large_scale_prediction.generator.utils import (
-    recover_global_position, unpad_global_coords)
+from aind_large_scale_prediction.generator.utils import recover_global_position, unpad_global_coords
 from aind_large_scale_prediction.io import ImageReaderFactory
 from cellpose import metrics
 from scipy.ndimage import binary_fill_holes, grey_dilation, map_coordinates
@@ -100,9 +99,7 @@ def create_initial_mask(
 
                 # Expand each voxel 3 voxels around it
                 for i, e in enumerate(expand):
-                    epix = (
-                        e[:, np.newaxis] + np.expand_dims(cell_centroids[k][i], 0) - 1
-                    )
+                    epix = e[:, np.newaxis] + np.expand_dims(cell_centroids[k][i], 0) - 1
                     # Flattenning points around a point inside ZYX
                     epix = epix.flatten()
 
@@ -155,7 +152,10 @@ def create_initial_mask(
 
 
 def remove_bad_flow_masks(
-    masks: ArrayLike, flows: ArrayLike, threshold: Optional[float] = 0.4, device=None
+    masks: ArrayLike,
+    flows: ArrayLike,
+    threshold: Optional[float] = 0.4,
+    device=None,
 ) -> ArrayLike:
     """
     Removes bad flows within the generated initial mask.
@@ -219,9 +219,7 @@ def fill_holes_and_remove_small_masks(
     """
 
     if masks.ndim > 3 or masks.ndim < 2:
-        raise ValueError(
-            "masks_to_outlines takes 2D or 3D array, not %dD array" % masks.ndim
-        )
+        raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" % masks.ndim)
 
     masks_properties = regionprops(masks)
 
@@ -388,9 +386,7 @@ def extract_global_to_local(
     ]
 
     # Mapping to the local coordinate system of the chunk
-    picked_global_ids_with_cells[..., :3] = (
-        picked_global_ids_with_cells[..., :3] - start_pos - pad
-    )
+    picked_global_ids_with_cells[..., :3] = picked_global_ids_with_cells[..., :3] - start_pos - pad
 
     # Validating seeds are within block boundaries
     picked_global_ids_with_cells = picked_global_ids_with_cells[
@@ -531,9 +527,7 @@ def execute_worker(
         f"Global slices: {global_coord_pos} - Unpadded global slices: {unpadded_global_slice[1:]} - Local slices: {unpadded_local_slice[1:]}"  # noqa: E501
     )
 
-    global_points_path = (
-        f"{cell_centroids_path}/global_seeds_{unpadded_global_slice[1:]}.npy"
-    )
+    global_points_path = f"{cell_centroids_path}/global_seeds_{unpadded_global_slice[1:]}.npy"
 
     # Unpadded block mask zeros if seeds don't exist in that area
     chunked_seg_mask = np.zeros(data.shape[1:], dtype=np.uint32)
@@ -685,6 +679,7 @@ def generate_masks(
 
         output_seg_dtype = np.uint32  # get_output_seg_data_type(n_cells=n_ids.shape[0])
         global_seeds = np.vstack((global_seeds.T, n_ids)).T
+        np.save(f"{results_folder}/cell_centroids.npy", global_seeds)
 
     else:
         raise ValueError("Please, provide the global seeds")
@@ -736,12 +731,8 @@ def generate_masks(
         multiprocessing.set_start_method("spawn", force=True)
 
     # Getting overlap prediction chunksize
-    overlap_prediction_chunksize = (0,) + tuple(
-        [axis_overlap * 2] * len(prediction_chunksize[-3:])
-    )
-    logger.info(
-        f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}"
-    )
+    overlap_prediction_chunksize = (0,) + tuple([axis_overlap * 2] * len(prediction_chunksize[-3:]))
+    logger.info(f"Overlap size based on cell diameter * 2: {overlap_prediction_chunksize}")
 
     lazy_data = (
         ImageReaderFactory()
@@ -799,18 +790,14 @@ def generate_masks(
 
     hists = zarr.open(hists_path, "r")
 
-    logger.info(
-        f"Creating masks in path: {output_seg_masks} chunks: {output_seg_masks.chunks}"
-    )
+    logger.info(f"Creating masks in path: {output_seg_masks} chunks: {output_seg_masks.chunks}")
 
     # Estimating total batches
     total_batches = np.prod(zarr_dataset.lazy_data.shape) / (
         np.prod(zarr_dataset.prediction_chunksize) * batch_size
     )
     samples_per_iter = n_workers * batch_size
-    logger.info(
-        f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}"
-    )
+    logger.info(f"Number of batches: {total_batches} - Samples per iteration: {samples_per_iter}")
 
     logger.info(f"{20*'='} Starting mask generation {20*'='}")
     start_time = time()

code/cell_segmentation/cellpose_segmentation/_shared/types.py code/aind_large_scale_cellpose/cellpose_segmentation/_shared/types.py

@@ -9,8 +9,7 @@
 import numpy as np
 from aind_large_scale_prediction._shared.types import ArrayLike
 from aind_large_scale_prediction.generator.utils import concatenate_lazy_data
-from aind_large_scale_prediction.generator.zarr_slice_generator import \
-    BlockedZarrArrayIterator
+from aind_large_scale_prediction.generator.zarr_slice_generator import BlockedZarrArrayIterator
 from aind_large_scale_prediction.io import extract_data
 from dask import config as da_cfg
 from distributed import Client, LocalCluster
@@ -72,9 +71,7 @@ def get_channel_percentiles(
         Dictionary with the computed percentiles.
     """
     # Iterate through the input array in steps equal to the block shape dimensions
-    slices_to_process = list(
-        BlockedZarrArrayIterator.gen_slices(array.shape, block_shape)
-    )
+    slices_to_process = list(BlockedZarrArrayIterator.gen_slices(array.shape, block_shape))
 
     percentiles = {}
 
@@ -206,17 +203,18 @@ def combine_percentiles(percentiles: Dict, method: Optional[str] = "min_max") ->
         channel_percentiles_cmb = None
         if method == "min_max":
             channel_percentiles_cmb = np.array(
-                [np.min(channel_percentiles[0]), np.max(channel_percentiles[1])]
+                [
+                    np.min(channel_percentiles[0]),
+                    np.max(channel_percentiles[1]),
+                ]
             )
 
         elif method == "median":
             channel_percentiles_cmb = np.median(channel_percentiles, axis=1)
 
         combined_percentiles.append(list(channel_percentiles_cmb))
 
-        print(
-            f"Channel {chn_idx}: {channel_percentiles} - cmb: {channel_percentiles_cmb}"
-        )
+        print(f"Channel {chn_idx}: {channel_percentiles} - cmb: {channel_percentiles_cmb}")
 
     return combined_percentiles
 
@@ -282,9 +280,7 @@ def compute_percentiles(
         threads_per_worker=threads_per_worker,
     )
 
-    combined_percentiles = combine_percentiles(
-        percentiles=percentiles, method=combine_method
-    )
+    combined_percentiles = combine_percentiles(percentiles=percentiles, method=combine_method)
 
     return combined_percentiles, percentiles