Image labeling at whole-well level

[tcompa]

As of #64 (comment), planned use cases for image labeling are

1. Segmentation at level 0 by site (original FOVs) in 2D & 3D
2. Segmentation for a full well, e.g. the full well in 2D (may still work at full resolution, i.e. pyramid level 0)
3. Segmentation at user-selected levels for 2D or 3D. e.g. Segment sites at pyramid level 1 or segment the whole well at level 2.

Use case 1 is discussed in #64.
This issue concerns use case 2, namely the possibility to segment a whole 2D well.

[tcompa]

As of 873e854, we have a first naive implementation, where the mask is explicitly upscaled before being stored on disk (see fractal-analytics-platform/fractal-tasks-core#24 for a discussion). This allowed for a quick implementation and visualization of results.

See for instance this result for a 2x2 well segmented at level 1:

Next steps:

1. Include the appropriate coordinateTransformations metadata, so that explicit upscaling is not needed for visualization (it will still be needed for processing, but probably done on-the-fly).
2. Integrate this task in the Fractal ones, to run it through Fractal/parsl.

Notes

• This task currently works on a single FOV, which corresponds to a well because we are always merging FOVs. In a more general case, a bit more of logic would be needed to load a whole well by stitching individual FOVs. Not needed, at the moment.
• Upscaling is done naively as in https://stackoverflow.com/a/7525345/19085332. Note: this is only tested on level 1 of a 2x2 well, but it may not scale up to a 9x8 well. To be checked. Anyway, upscaling won't be part of this task (see point 1 below), so this discussion can be postponed to Create core feature measurements #69.

[jluethi]

This segmentation result looks great @tcompa ! 👏🏻

In a more general case, a bit more of logic would be needed to load a whole well by stitching individual FOVs. Not needed, at the moment.

With our current approach being wells always being saved as single FOVs, we don't need to go into this now. But worth considering down the road :)

Very much looking forward to seeing this run on whole wells through Fractal. I will be preparing additional test sets where we want to test 2D segmentation of wells on organoids (much larger structures), which should work by changing some of the cellpose parameters :)

[tcompa]

As a future reference, segmenting an array of shape (1, 9720, 10240) (corresponding to level 1 of a 2D 9x8 well) leads to a GPU-memory error within cellpose:

Traceback (most recent call last):
  File "../../image_labeling_whole_well.py", line 210, in <module>
    image_labeling_whole_well(
  File "../../image_labeling_whole_well.py", line 101, in image_labeling_whole_well
    mask, flows, styles, diams = model.eval(
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/models.py", line 227, in eval
    masks, flows, styles = self.cp.eval(x, 
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/models.py", line 536, in eval
    masks, styles, dP, cellprob, p = self._run_cp(x, 
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/models.py", line 636, in _run_cp
    outputs = dynamics.compute_masks(dP[:,i], cellprob[i], niter=niter, cellprob_threshold=cellprob_threshold,
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/dynamics.py", line 725, in compute_masks
    mask = remove_bad_flow_masks(mask, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device)
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/dynamics.py", line 580, in remove_bad_flow_masks
    merrors, _ = metrics.flow_error(masks, flows, use_gpu, device)
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/metrics.py", line 282, in flow_error
    dP_masks = dynamics.masks_to_flows(maski, use_gpu=use_gpu, device=device)
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/dynamics.py", line 277, in masks_to_flows
    mu, mu_c = masks_to_flows_device(masks, device=device)
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/dynamics.py", line 152, in masks_to_flows_gpu
    mu = _extend_centers_gpu(neighbors, centers, isneighbor, Ly, Lx, 
  File "/data/homes/fractal/.conda/envs/fractal/lib/python3.8/site-packages/cellpose/dynamics.py", line 83, in _extend_centers_gpu
    grads = T[:, pt[[2,1,4,3],:,0], pt[[2,1,4,3],:,1]]
RuntimeError: CUDA out of memory. Tried to allocate 1.29 GiB (GPU 0; 14.76 GiB total capacity; 12.42 GiB already allocated; 561.44 MiB free; 13.29 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Now trying with level 2.

[tcompa]

Labeling a single Z plane of a 9x8 well at level 2 (corresponding to a x4 coarsening along X and Y) does not lead to memory errors:

Start image_labeling_whole_well task for /data/active/fractal/tests/Temporary_data_UZH_1_well_9x8_sites_singlefov/20200812-CardiomyocyteDifferentiation14-Cycle1_mip.zarr/B/03/0/
Total well shape/chunks:
(1, 4860, 5120)
((1,), (2160, 2160, 540), (2560, 2560))

START Cellpose | well: <class 'dask.array.core.Array'>, (1, 4860, 5120) 
END   Cellpose | Elapsed: 19.3520 seconds | mask shape: (1, 4860, 5120), mask dtype: uint16, max(mask): 8380

upscaled mask from (1, 4860, 5120) to (1, 19440, 20480)

da.from_array(upscaled_mask) [with rechunking]: dask.array<rechunk-merge, shape=(1, 19440, 20480), dtype=uint16, chunksize=(1, 2160, 2560), chunktype=numpy.ndarray>

The memory peak for this run is ~7 G, and the runtime is ~20 seconds.

Labeling options are probably somewhat off, since it seems that we find less labels than expected:

[tcompa]

BUGFIX:

In the previous comment I kept diameter=40.0, while in fact it needs to change when working on higher levels (since it is in pixel units). After setting diameter=(40.0 / coarsening_xy ** labeling_level) (see bae9bd8), segmentation takes ~1 minute, with peak memory at ~9 G.

And the labels look much more reasonable:

[jluethi]

@tcompa This looks great! :)
Yes, adapting the diameter certainly is necessary here. It describes how many pixels an object is approximately taking over.

diameter=(40.0 / coarsening_xy ** labeling_level)

This is a good proxy given the input 40. May even be an interesting setup for the users afterwards: We take the diameter at full res and calculate a relevant diameter parameter for the pyramid level

[tcompa]

As a future reference, when labeling (e.g.) level 2 (and with coarsening_xy=2), the correct .zattrs file for the labels should include

            "datasets": [
                {
                    "path": "0",
                    "coordinateTransformations": [{
                        "type": "scale",
                        "scale": [1.0, 4.0, 4.0]
                    }]
                },
                ...
            ],

I'm still trying to understand what scale arguments should go into the next levels ("path": "1", "path": "2", ..). My doubt is: should scale remain fixed and equal to [1,4,4] for all paths, or should it change as [1,4,4], [1,8,8], [1,16,16], ...?

This is hard to understand ex-post, because it seems that napari doesn't load pyramids when showing a single FOV - and I cannot show a well/plate, since it seems that they do not support labels (*). Thus I'm only viewing level 0.

@jluethi, is this question possibly something obvious to you?

(*) See ongoing work on ome/ome-zarr-py#207 and ome/ome-zarr-py#65.