Support 2D-only OME-Zarrs?

[jluethi]

We save our 2D OME-Zarrs as (z, y, x) images with (1, size_y, size_x). Some other OME-Zarrs (e.g. the ones generated by the faim-hcs OME-Zarr converter) can be (x,y) only, i.e. (size_y, size_x).

This then leads to issues in our processing, for example the following error in cellpose. The problem here is a mismatch between ROI dimensions (ROIs always being z, y, x) and image dimensions.

2023-06-07 13:03:23,100; INFO; START cellpose_segmentation task
2023-06-07 13:03:23,100; INFO; zarrurl='/Users/joel/Desktop/20230607_MD_Maurice_test2/Plate.zarr/C/3/0'
2023-06-07 13:03:23,217; INFO; data_zyx.shape=(3072, 3072)
2023-06-07 13:03:23,242; INFO; ROI table at /Users/joel/Desktop/20230607_MD_Maurice_test2/Plate.zarr/C/3/0/tables/well_ROI_table has attrs: <zarr.attrs.Attributes object at 0x29e828580>
2023-06-07 13:03:23,242; INFO; ROI table at /Users/joel/Desktop/20230607_MD_Maurice_test2/Plate.zarr/C/3/0/tables/well_ROI_table cannot be used for masked loading. Set use_masks=False.
2023-06-07 13:03:23,242; INFO; use_masks=False
2023-06-07 13:03:23,242; INFO; reset_origin=True
2023-06-07 13:03:23,243; INFO; pxl_zyx=[1.0, 0.6834, 0.6834]
2023-06-07 13:03:23,243; INFO; new_labels=['organoids']
2023-06-07 13:03:23,243; INFO; existing_labels=[]
2023-06-07 13:03:23,251; INFO; Output label path: /Users/joel/Desktop/20230607_MD_Maurice_test2/Plate.zarr/C/3/0/labels/organoids/0
2023-06-07 13:03:23,261; INFO; mask will have shape (3072, 3072) and chunks ((2048, 1024), (2048, 1024))
2023-06-07 13:03:23,263; INFO; TORCH CUDA version not installed/working.
2023-06-07 13:03:23,265; INFO; >> cyto2 << model set to be used
2023-06-07 13:03:23,265; INFO; >>>> using CPU
2023-06-07 13:03:23,265; INFO; WARNING: MKL version on torch not working/installed - CPU version will be slightly slower.
2023-06-07 13:03:23,265; INFO; see https://pytorch.org/docs/stable/backends.html?highlight=mkl
2023-06-07 13:03:23,369; INFO; >>>> model diam_mean =  30.000 (ROIs rescaled to this size during training)
2023-06-07 13:03:23,369; INFO; Start cellpose_segmentation task for /Users/joel/Desktop/20230607_MD_Maurice_test2/Plate.zarr/C/3/0
2023-06-07 13:03:23,369; INFO; relabeling: True
2023-06-07 13:03:23,369; INFO; do_3D: True
2023-06-07 13:03:23,369; INFO; use_gpu: False
2023-06-07 13:03:23,369; INFO; level: 1
2023-06-07 13:03:23,369; INFO; model_type: cyto2
2023-06-07 13:03:23,369; INFO; pretrained_model: None
2023-06-07 13:03:23,369; INFO; anisotropy: 1.4632718759145449
2023-06-07 13:03:23,369; INFO; Total well shape/chunks:
2023-06-07 13:03:23,369; INFO; (3072, 3072)
2023-06-07 13:03:23,369; INFO; ((2048, 1024), (2048, 1024))
2023-06-07 13:03:23,369; INFO; Now starting loop over 1 ROIs
2023-06-07 13:03:23,369; INFO; Now processing ROI 1/1
Traceback (most recent call last):
  File "/Users/joel/Dropbox/Joel/FMI/Code/fractal/fractal-demos/examples/server/FRACTAL_TASKS_DIR/.fractal/fractal-tasks-core0.9.4/venv/lib/python3.9/site-packages/fractal_tasks_core/cellpose_segmentation.py", line 699, in <module>
    run_fractal_task(
  File "/Users/joel/Dropbox/Joel/FMI/Code/fractal/fractal-demos/examples/server/FRACTAL_TASKS_DIR/.fractal/fractal-tasks-core0.9.4/venv/lib/python3.9/site-packages/fractal_tasks_core/_utils.py", line 91, in run_fractal_task
    metadata_update = task_function(**task_args.dict(exclude_unset=True))
  File "/Users/joel/Dropbox/Joel/FMI/Code/fractal/fractal-demos/examples/server/FRACTAL_TASKS_DIR/.fractal/fractal-tasks-core0.9.4/venv/lib/python3.9/site-packages/fractal_tasks_core/cellpose_segmentation.py", line 516, in cellpose_segmentation
    img_np = np.expand_dims(data_zyx[region].compute(), axis=0)
  File "/Users/joel/Dropbox/Joel/FMI/Code/fractal/fractal-demos/examples/server/FRACTAL_TASKS_DIR/.fractal/fractal-tasks-core0.9.4/venv/lib/python3.9/site-packages/dask/array/core.py", line 1978, in __getitem__
    index2 = normalize_index(index, self.shape)
  File "/Users/joel/Dropbox/Joel/FMI/Code/fractal/fractal-demos/examples/server/FRACTAL_TASKS_DIR/.fractal/fractal-tasks-core0.9.4/venv/lib/python3.9/site-packages/dask/array/slicing.py", line 904, in normalize_index
    raise IndexError("Too many indices for array")
IndexError: Too many indices for array

TODO: Figure out whether we can support 2D OME-Zarr images 1) by adapting the loading to support that? 2) by adapting our ROI tables for 2D images?

I'd prefer 1, but need to evaluate

[jluethi]

This is actually an example of hitting this issue here: https://github.com/fractal-analytics-platform/fractal-tasks-core/issues/150 Specifically point 3: Make sure that the relevant functions/tasks are capable of handling arrays of different shapes.

The issue arises primarily because we have 3D ROIs & 2D images. But even if our ROIs would be 2D as well (=> the indexing would work if we also handle that case), the cellpose function would still always return a 3D output image (the way it's currently programmed) and we have some other 3D assumptions in its code.

2 takeaways: 1) We should pursue the flexible dimensionality handling (i.e. I think illumination correction will potentially also run into some issues here on actual 2D inputs where images should be overwritten) 2) I want to test whether we can store our resulting labels in our current 3D mode (if yes, then we just need to handle loading of 2D data correctly).

---

To attempt 2: I think we can make this work by adding a check during ROI loading. That would probably mean having a library function for ROI loading which would handle the 2D case differently (checking for it and still returning a 3D array with dim0 = 1). I'll investigate whether this is works and whether we have downstream issues when saving the labels in this 3D mode (I think not, but maybe there would be issues with the viewer?)

[tcompa]

How would one of the 2D-only input OME-Zarrs that we'd like to support look like? Would it be fully CYX like the one below?

I'm trying to write a test with this kind of input, but I'm not sure this was your intended goal @jluethi, nor I'm sure that #403 is enough (I'm still reviewing..).

        'multiscales': [
            {
                'axes': [                         ############### NOTE: There is no Z
                    {
                        'name': 'c',
                        'type': 'channel',
                    },
                    {
                        'name': 'y',
                        'type': 'space',
                        'unit': 'micrometer',
                    },
                    {
                        'name': 'x',
                        'type': 'space',
                        'unit': 'micrometer',
                    },
                ],
                'datasets': [
                    {
                        'coordinateTransformations': [
                            {
                                'scale': [         ############## NOTE: CYX
                                    1.0,
                                    0.1625,
                                    0.1625,
                                ],
                                'type': 'scale',
                            },
                        ],
                        'path': '0',
                    },
                    {
                        'coordinateTransformations': [
                            {
                                'scale': [
                                    1.0,
                                    0.325,
                                    0.325,
                                ],
                                'type': 'scale',
                            },
                        ],
                        'path': '1',
                    },
                    {
                        'coordinateTransformations': [
                            {
                                'scale': [
                                    1.0,
                                    0.65,
                                    0.65,
                                ],
                                'type': 'scale',
                            },
                        ],
                        'path': '2',
                    },
                    {
                        'coordinateTransformations': [
                            {
                                'scale': [
                                    1.0,
                                    1.3,
                                    1.3,
                                ],
                                'type': 'scale',
                            },
                        ],
                        'path': '3',
                    },
                    {
                        'coordinateTransformations': [
                            {
                                'scale': [
                                    1.0,
                                    2.6,
                                    2.6,
                                ],
                                'type': 'scale',
                            },
                        ],
                        'path': '4',
                    },
                ],
                'version': '0.4',
            },
        ],
        'omero': {
            'channels': [
                {
                    'active': True,
                    'coefficient': 1,
                    'color': '00FFFF',
                    'family': 'linear',
                    'inverted': False,
                    'label': 'DAPI',
                    'wavelength_id': 'A01_C01',
                    'window': {
                        'end': 800,
                        'max': 65535,
                        'min': 0,
                        'start': 110,
                    },
                },
            ],
            'id': 1,
            'name': 'TBD',
            'version': '0.4',
        },
    }

[tcompa]

How would one of the 2D-only input OME-Zarrs that we'd like to support look like? Would it be fully CYX like the one below? I'm trying to write a test with this kind of input, but I'm not sure this was your intended goal @jluethi, nor I'm sure that https://github.com/fractal-analytics-platform/fractal-tasks-core/pull/403 is enough (I'm still reviewing..).

I'm still not convinced that #403 would not work for this kind of OME-Zarr input. In fact it may be that it does, by accident, by interpreting the C dimension as a Z dimension - when parsing the ZYX pixel sizes (via extract_zyx_pixel_sizes). But this only works if the scale transformation has value 1 for both the C and Z axis (which is not necessarily true).

[tcompa]

With https://github.com/fractal-analytics-platform/fractal-tasks-core/commit/bad56c4cbe82177f392dd6bc9a1f389963f47f3c, I added a test that runs the Cellpose task on a true CZY OME-Zarr array. The test actually passes, and it produces labels, but I think this is again not working correctly. This is clear by looking at the logs

INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:296 zarrurl='/tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0'
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:364 data_zyx.shape=(540, 1280)
INFO     fractal_tasks_core.lib_regions_of_interest:lib_regions_of_interest.py:400 ROI table at /tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0/tables/FOV_ROI_table has attrs: {'encoding-type': 'anndata', 'encoding-version': '0.1.0'}
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:378 ROI table at /tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0/tables/FOV_ROI_table cannot be used for masked loading. Set use_masks=False.
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:383 use_masks=False
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:389 full_res_pxl_sizes_zyx=[1.0, 0.1625, 0.1625]
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:393 actual_res_pxl_sizes_zyx=[1.0, 0.65, 0.65]
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:400 reset_origin=True
INFO     fractal_tasks_core.lib_regions_of_interest:lib_regions_of_interest.py:226 full_res_pxl_sizes_zyx=[1.0, 0.1625, 0.1625]
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:471 new_labels=['label_DAPI']
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:472 existing_labels=[]
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:496 Output label path: /tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0/labels/label_DAPI/0
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:517 mask will have shape (540, 1280) and chunks ((540,), (1280,))
test_workflows_cellpose_segmentation.py:161 patched_cellpose_core_use_gpu
    'WARNING: using patched_cellpose_core_use_gpu' (str) len=44
INFO     cellpose.models:models.py:338 >> cyto2 << model set to be used
INFO     cellpose.core:core.py:90 >>>> using CPU
INFO     cellpose.models:models.py:372 >>>> model diam_mean =  30.000 (ROIs rescaled to this size during training)
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:532 Start cellpose_segmentation task for /tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:533 relabeling: True
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:534 do_3D: False
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:535 use_gpu: False
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:536 level: 2
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:537 model_type: cyto2
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:538 pretrained_model: None
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:539 anisotropy: None
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:540 Total well shape/chunks:
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:541 (540, 1280)
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:542 ((540,), (1280,))
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:558 Now starting loop over 2 ROIs
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:567 Now processing ROI 1/2
INFO     cellpose_segmentation.py:test_workflows_cellpose_segmentation.py:119 [None][patched_segment_ROI] START
INFO     cellpose_segmentation.py:test_workflows_cellpose_segmentation.py:132 [None][patched_segment_ROI] END
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:636 ROI [0, 1, 0, 540, 0, 640], num_labels_roi=2, num_labels_tot=2
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:567 Now processing ROI 2/2
INFO     cellpose_segmentation.py:test_workflows_cellpose_segmentation.py:119 [None][patched_segment_ROI] START
INFO     cellpose_segmentation.py:test_workflows_cellpose_segmentation.py:132 [None][patched_segment_ROI] END
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:636 ROI [0, 1, 0, 540, 640, 1280], num_labels_roi=2, num_labels_tot=4
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:672 End cellpose_segmentation task for /tmp/pytest-of-tommaso/pytest-93/test_CYX_input0/tmp_out_mip/plate_mip.zarr/B/03/0, now building pyramids.
INFO     fractal_tasks_core.tasks.cellpose_segmentation:cellpose_segmentation.py:688 End building pyramids

where it's clear that the 1.0 which is the "pixel size" along the C axis is interpreted as the one along the Z axis.

[tcompa]

As of our discussion a few minutes ago:

• [x] We are not really supporting the case where the Z axis is missing. It may work, because its "scale" is typically the same as the C one (i.e. 1.0), but we should at least have a warning.
• [x] We are currently always relying that the first axis is C, independently on what comes next. Because of this, we can simply fix a current issue with the labeling tasks (they produce wrong metadata, which has scale transformations with four elements) in a simple way: always strip away the first item, since labels do not have C axis.
• [x] What about MIP? Does it also rely on having three-items scale transformations? Let's review this.
• [x] We are clearly not fixing #150 here, but let's aim at a version where example 01 in fractal-demos does work.
• [x] (for me) I've reviewed the cellpose task in #403, and I still need to review the napari-workflows one.
• [x] Add a CYX test.

[jluethi]

Great overview of conclusions from the call, this should be the scope limit for merging it.

Would it be fully CYX like the one below?

Yes, the examples I have so far always are like this and I'm aiming for supporting CYX support, not YX support. Yes, that this works is somewhat accidental, we'll need to improve with #150 eventually.

We are currently always relying that the first axis is C, independently on what comes next. Because of this, we can simply fix a current issue with the labeling tasks (they produce wrong metadata, which has scale transformations with four elements) in a simple way: always strip away the first item, since labels do not have C axis.

If we include a check for the presence of a C axis in relevant library functions, then that should be save and show us where we'd need to continue when we get back to it.

[tcompa]

• [x] Add a check that the first axis is c, in some relevant library function - which one(s)?

[tcompa]

The new test_workflows_napari_workflows.py::test_napari_workflow_CYX (90c7528) applies two napari workflows (labeling and measurement) to a CYX OME-Zarr input. The final OME-Zarr

• has three-dimensional (that is, CYX) metadata both for the image and for the labels
• has three-dimensional (that is, CYX) structure of the zarr arrays (that is, in .zattrs)
• renders correctly in napari.

[jluethi]

That's great! 👏🏻

[tcompa]

We are clearly not fixing https://github.com/fractal-analytics-platform/fractal-tasks-core/issues/150 here, but let's aim at a version where example 01 in fractal-demos does work.

I've not run example 01 in fractal-demos (as of https://github.com/fractal-analytics-platform/fractal-demos/pull/54/commits/1152116bcd6e1c20d2960915e892b1d3cef8e096), and it went through with no noticeable issues.

• The 3D OME-Zarr is present, and can be viewed in napari
• The 2D OME-Zarr is present, and can be viewed in napari (including its labels)
• Both have CZYX axes, and in both cases the scale transformations have four items
• The label image of the MIP OME-Zarr has ZYX axes and scale transformations with three items

[jluethi]

That's great! Sounds like we're ready to merge this then.

[tcompa]

Add a check that the first axis is c, in some relevant library function - which one(s)?

See https://github.com/fractal-analytics-platform/fractal-tasks-core/pull/403/commits/09f0c8dd14f1f2c8afab1b92995be204a6e92900

[jluethi]

I also ran example 01 locally now and all looks good to me there. I'll continue work to update the MD converter task for server 1.3.0 & tasks 0.10.0 compatibility and will test that afterwards, but would open new issues and a new PR if problems come up there.

[tcompa]

What about MIP? Does it also rely on having three-items scale transformations? Let's review this.

When we have a CZYX OME-Zarr and we call copy_ome_zarr.py with project_to_2D=True, the output is a new (metadata-only) CZYX OME-Zarr, which simply has a dummy Z axis.

Both copy_ome_zarr.py and maximum_intensity_projection.py then use extract_zyx_pixel_sizes, which we patched so that it alwasy returns the last three elements of scale transformations as ZYX pixel sizes.

TL;DR Even if we don't have much flexibility, in the current (transitional..) state of fractal-tasks-core, I don't see any special issue with MIP in current workflows:

• A Fractal-generated 3D zarr should always be a CZYX one;
• An CYX existing OME-Zarr would be out-of-scope for the MIP task;
• A ZYX existing OME-Zarr would be in-scope for the MIP task, but only as part of #150.

[tcompa]

That's great! Sounds like we're ready to merge this then.

There's nothing left on my side re: #403.