From 364c5e3a8ca4f8f586a045283d0d5cb8db32cc84 Mon Sep 17 00:00:00 2001
From: jluethi <joel.luethi@uzh.ch>
Date: Tue, 16 Jul 2024 11:37:58 +0200
Subject: [PATCH] Fix for yx images in Cellpose task

---
 CHANGELOG.md                                  |  1 +
 .../tasks/cellpose_segmentation.py            | 41 ++++++++++++-------
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aef5d39a4..883895907 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 * Tasks:
     * Fix issue with masked ROI & relabeling in Cellpose task (\#785).
     * Fix issue with masking ROI label types in masked_loading_wrapper for Cellpose task (\#785).
+    * Enable workaround to support yx images in Cellpose task (\#789).
 
 # 1.1.0
 
diff --git a/fractal_tasks_core/tasks/cellpose_segmentation.py b/fractal_tasks_core/tasks/cellpose_segmentation.py
index 13cb8fe1a..d371aea24 100644
--- a/fractal_tasks_core/tasks/cellpose_segmentation.py
+++ b/fractal_tasks_core/tasks/cellpose_segmentation.py
@@ -314,11 +314,21 @@ def cellpose_segmentation(
             output_label_name = f"label_{ind_channel}"
 
     # Load ZYX data
-    data_zyx = da.from_zarr(f"{zarr_url}/{level}")[ind_channel]
+    # Workaround for #788: Only load channel index when there is a channel
+    # dimension
+    if ngff_image_meta.axes_names[0] != "c":
+        data_zyx = da.from_zarr(f"{zarr_url}/{level}")
+        if channel2.is_set():
+            raise ValueError(
+                "Dual channel input was specified for an OME-Zarr image "
+                "without a channel axis"
+            )
+    else:
+        data_zyx = da.from_zarr(f"{zarr_url}/{level}")[ind_channel]
+        if channel2.is_set():
+            data_zyx_c2 = da.from_zarr(f"{zarr_url}/{level}")[ind_channel_c2]
+            logger.info(f"Second channel: {data_zyx_c2.shape=}")
     logger.info(f"{data_zyx.shape=}")
-    if channel2.is_set():
-        data_zyx_c2 = da.from_zarr(f"{zarr_url}/{level}")[ind_channel_c2]
-        logger.info(f"Second channel: {data_zyx_c2.shape=}")
 
     # Read ROI table
     ROI_table_path = f"{zarr_url}/tables/{input_ROI_table}"
@@ -365,18 +375,21 @@ def cellpose_segmentation(
         logger.info(f"Anisotropy: {advanced_cellpose_model_params.anisotropy}")
 
     # Rescale datasets (only relevant for level>0)
+    # Workaround for #788
     if ngff_image_meta.axes_names[0] != "c":
-        raise ValueError(
-            "Cannot set `remove_channel_axis=True` for multiscale "
-            f"metadata with axes={ngff_image_meta.axes_names}. "
-            'First axis should have name "c".'
+        new_datasets = rescale_datasets(
+            datasets=[ds.dict() for ds in ngff_image_meta.datasets],
+            coarsening_xy=coarsening_xy,
+            reference_level=level,
+            remove_channel_axis=False,
+        )
+    else:
+        new_datasets = rescale_datasets(
+            datasets=[ds.dict() for ds in ngff_image_meta.datasets],
+            coarsening_xy=coarsening_xy,
+            reference_level=level,
+            remove_channel_axis=True,
         )
-    new_datasets = rescale_datasets(
-        datasets=[ds.dict() for ds in ngff_image_meta.datasets],
-        coarsening_xy=coarsening_xy,
-        reference_level=level,
-        remove_channel_axis=True,
-    )
 
     label_attrs = {
         "image-label": {