Merge pull request #709 from fractal-analytics-platform/703_fix_illum…

…_corr_no_overwrite

703 fix illum corr no overwrite

fractal_tasks_core/tasks/_registration_utils.py

@@ -12,16 +12,7 @@
 import pandas as pd
 
 from fractal_tasks_core.ngff.zarr_utils import load_NgffWellMeta
-
-
-def _split_well_path_image_path(zarr_url: str) -> tuple[str, str]:
-    """
-    Returns path to well folder for HCS OME-Zarr `zarr_url`.
-    """
-    zarr_url = zarr_url.rstrip("/")
-    well_path = "/".join(zarr_url.split("/")[:-1])
-    img_path = zarr_url.split("/")[-1]
-    return well_path, img_path
+from fractal_tasks_core.tasks._zarr_utils import _split_well_path_image_path
 
 
 def create_well_acquisition_dict(

fractal_tasks_core/tasks/_zarr_utils.py

@@ -0,0 +1,104 @@
+import copy
+
+import zarr
+from filelock import FileLock
+
+from fractal_tasks_core.ngff.zarr_utils import load_NgffWellMeta
+
+
+def _copy_hcs_ome_zarr_metadata(
+    zarr_url_origin: str,
+    zarr_url_new: str,
+) -> None:
+    """
+    Updates the necessary metadata for a new copy of an OME-Zarr image
+
+    Based on an existing OME-Zarr image in the same well, the metadata is
+    copied and added to the new zarr well. Additionally, the well-level
+    metadata is updated to include this new image.
+
+    Args:
+        zarr_url_origin: zarr_url of the origin image
+        zarr_url_new: zarr_url of the newly created image. The zarr-group
+            already needs to exist, but metadata is written by this function.
+    """
+    # Copy over OME-Zarr metadata for illumination_corrected image
+    # See #681 for discussion for validation of this zattrs
+    old_image_group = zarr.open_group(zarr_url_origin, mode="r")
+    old_attrs = old_image_group.attrs.asdict()
+    zarr_url_new = zarr_url_new.rstrip("/")
+    new_image_group = zarr.group(zarr_url_new)
+    new_image_group.attrs.put(old_attrs)
+
+    # Update well metadata about adding the new image:
+    new_image_path = zarr_url_new.split("/")[-1]
+    well_url, old_image_path = _split_well_path_image_path(zarr_url_origin)
+    _update_well_metadata(well_url, old_image_path, new_image_path)
+
+
+def _update_well_metadata(
+    well_url: str,
+    old_image_path: str,
+    new_image_path: str,
+    timeout: int = 120,
+) -> None:
+    """
+    Update the well metadata by adding the new_image_path to the image list.
+
+    The content of new_image_path will be based on old_image_path, the origin
+    for the new image that was created.
+    This function aims to avoid race conditions with other processes that try
+    to update the well metadata file by using FileLock & Timeouts
+
+    Args:
+        well_url: Path to the HCS OME-Zarr well that needs to be updated
+        old_image_path: path relative to well_url where the original image is
+            found
+        new_image_path: path relative to well_url where the new image is placed
+        timeout: Timeout in seconds for trying to get the file lock
+    """
+    lock = FileLock(f"{well_url}/.zattrs.lock")
+    with lock.acquire(timeout=timeout):
+
+        well_meta = load_NgffWellMeta(well_url)
+        existing_well_images = [image.path for image in well_meta.well.images]
+        if new_image_path in existing_well_images:
+            raise ValueError(
+                f"Could not add the {new_image_path=} image to the well "
+                "metadata because and image with that name "
+                f"already existed in the well metadata: {well_meta}"
+            )
+        try:
+            well_meta_image_old = next(
+                image
+                for image in well_meta.well.images
+                if image.path == old_image_path
+            )
+        except StopIteration:
+            raise ValueError(
+                f"Could not find an image with {old_image_path=} in the "
+                "current well metadata."
+            )
+        well_meta_image = copy.deepcopy(well_meta_image_old)
+        well_meta_image.path = new_image_path
+        well_meta.well.images.append(well_meta_image)
+        well_meta.well.images = sorted(
+            well_meta.well.images,
+            key=lambda _image: _image.path,
+        )
+
+        well_group = zarr.group(well_url)
+        well_group.attrs.put(well_meta.dict(exclude_none=True))
+
+    # One could catch the timeout with a try except Timeout. But what to do
+    # with it?
+
+
+def _split_well_path_image_path(zarr_url: str) -> tuple[str, str]:
+    """
+    Returns path to well folder for HCS OME-Zarr `zarr_url`.
+    """
+    zarr_url = zarr_url.rstrip("/")
+    well_path = "/".join(zarr_url.split("/")[:-1])
+    img_path = zarr_url.split("/")[-1]
+    return well_path, img_path

fractal_tasks_core/tasks/apply_registration_to_image.py

@@ -35,7 +35,7 @@
 from fractal_tasks_core.roi import is_standard_roi_table
 from fractal_tasks_core.roi import load_region
 from fractal_tasks_core.tables import write_table
-from fractal_tasks_core.tasks._registration_utils import (
+from fractal_tasks_core.tasks._zarr_utils import (
     _split_well_path_image_path,
 )
 from fractal_tasks_core.utils import _get_table_path_dict

fractal_tasks_core/tasks/illumination_correction.py

@@ -33,6 +33,7 @@
 from fractal_tasks_core.roi import (
     convert_ROI_table_to_indices,
 )
+from fractal_tasks_core.tasks._zarr_utils import _copy_hcs_ome_zarr_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -142,9 +143,9 @@ def illumination_correction(
 
     # Defione old/new zarrurls
     if overwrite_input:
-        zarr_url_new = zarr_url
+        zarr_url_new = zarr_url.rstrip("/")
     else:
-        zarr_url_new = zarr_url + suffix
+        zarr_url_new = zarr_url.rstrip("/") + suffix
 
     t_start = time.perf_counter()
     logger.info("Start illumination_correction")
@@ -226,6 +227,7 @@ def illumination_correction(
             overwrite=False,
             dimension_separator="/",
         )
+        _copy_hcs_ome_zarr_metadata(zarr_url, zarr_url_new)
 
     # Iterate over FOV ROIs
     num_ROIs = len(list_indices)

tests/data/generate_zarr_ones.py

@@ -28,6 +28,23 @@
 
 if os.path.isdir(zarrurl):
     shutil.rmtree(zarrurl)
+
+plate_group = zarr.open_group(zarrurl)
+plate_group.attrs.put(
+    {
+        "plate": {
+            "acquisitions": [{"id": 1, "name": "plate_ones"}],
+            "columns": [{"name": "03"}],
+            "rows": [{"name": "B"}],
+            "version": "0.4",
+            "wells": [{"columnIndex": 0, "path": "B/03", "rowIndex": 0}],
+        }
+    }
+)
+
+well_group = zarr.open(f"{zarrurl}B/03/")
+well_group.attrs.put({"well": {"images": [{"path": "0"}], "version": "0.4"}})
+
 component = "B/03/0/"
 
 for ind_level in range(num_levels):

tests/data/plate_ones.zarr/.zattrs

@@ -0,0 +1,28 @@
+{
+    "plate": {
+        "acquisitions": [
+            {
+                "id": 1,
+                "name": "plate_ones"
+            }
+        ],
+        "columns": [
+            {
+                "name": "03"
+            }
+        ],
+        "rows": [
+            {
+                "name": "B"
+            }
+        ],
+        "version": "0.4",
+        "wells": [
+            {
+                "columnIndex": 0,
+                "path": "B/03",
+                "rowIndex": 0
+            }
+        ]
+    }
+}

tests/data/plate_ones.zarr/B/03/.zattrs

@@ -0,0 +1,10 @@
+{
+    "well": {
+        "images": [
+            {
+                "path": "0"
+            }
+        ],
+        "version": "0.4"
+    }
+}

tests/tasks/test_unit_illumination_correction.py

@@ -6,24 +6,31 @@
 import anndata as ad
 import dask.array as da
 import numpy as np
+import pytest
 from pytest import LogCaptureFixture
 from pytest import MonkeyPatch
 
 from fractal_tasks_core.ngff.zarr_utils import load_NgffImageMeta
+from fractal_tasks_core.ngff.zarr_utils import load_NgffWellMeta
 from fractal_tasks_core.roi import (
     convert_ROI_table_to_indices,
 )
+from fractal_tasks_core.tasks._registration_utils import (
+    _split_well_path_image_path,
+)
 from fractal_tasks_core.tasks.illumination_correction import correct
 from fractal_tasks_core.tasks.illumination_correction import (
     illumination_correction,
 )
 
 
+@pytest.mark.parametrize("overwrite_input", [True, False])
 def test_illumination_correction(
     tmp_path: Path,
     testdata_path: Path,
     monkeypatch: MonkeyPatch,
     caplog: LogCaptureFixture,
+    overwrite_input: bool,
 ):
     # GIVEN a zarr pyramid on disk, made of all ones
     # WHEN I apply illumination_correction
@@ -49,13 +56,9 @@ def test_illumination_correction(
     with open(zarr_url + ".zattrs") as fin:
         zattrs = json.load(fin)
         num_levels = len(zattrs["multiscales"][0]["datasets"])
-    metadata: dict = {
-        "num_levels": num_levels,
-        "coarsening_xy": 2,
-    }
-    num_channels = 2
-    num_levels = metadata["num_levels"]
 
+    num_channels = 2
+    num_levels = num_levels
     # Read FOV ROIs and create corresponding indices
     ngff_image_meta = load_NgffImageMeta(zarr_url)
     pixels = ngff_image_meta.get_pixel_sizes_zyx(level=0)
@@ -82,14 +85,15 @@ def patched_correct(*args, **kwargs):
         "fractal_tasks_core.tasks.illumination_correction.correct",
         patched_correct,
     )
-
+    suffix = "_illum_corr"
     # Call illumination correction task, with patched correct()
     illumination_correction(
         zarr_url=zarr_url,
-        overwrite_input=True,
+        overwrite_input=overwrite_input,
         illumination_profiles_folder=illumination_profiles_folder,
         illumination_profiles=illum_params,
         background=0,
+        suffix=suffix,
     )
 
     print(caplog.text)
@@ -100,13 +104,35 @@ def patched_correct(*args, **kwargs):
         tot_calls_correct = len(f.read().splitlines())
     assert tot_calls_correct == expected_tot_calls_correct
 
+    old_urls = [testdata_path / "plate_ones.zarr/B/03/0"]
+    if overwrite_input:
+        new_zarr_url = zarr_url.rstrip("/")
+    else:
+        new_zarr_url = zarr_url.rstrip("/") + suffix
+        old_urls.append(zarr_url.rstrip("/"))
+
     # Verify the output
-    for ind_level in range(num_levels):
-        old = da.from_zarr(
-            testdata_path / f"plate_ones.zarr/B/03/0/{ind_level}"
-        )
-        new = da.from_zarr(f"{zarr_url}{ind_level}")
-        assert old.shape == new.shape
-        assert old.chunks == new.chunks
-        assert new.compute()[0, 0, 0, 0] == 1
-        assert np.allclose(old.compute(), new.compute())
+    for old_url in old_urls:
+        for ind_level in range(num_levels):
+            old = da.from_zarr(f"{old_url}/{ind_level}")
+            print(testdata_path / f"plate_ones.zarr/B/03/0/{ind_level}")
+            print(f"{zarr_url}{ind_level}")
+            new = da.from_zarr(f"{new_zarr_url}/{ind_level}")
+            assert old.shape == new.shape
+            assert old.chunks == new.chunks
+            assert new.compute()[0, 0, 0, 0] == 1
+            assert np.allclose(old.compute(), new.compute())
+
+    # Verify that the new_zarr_url has valid OME-Zarr metadata
+    _ = load_NgffImageMeta(new_zarr_url)
+
+    # Verify the well metadata: Are all the images in well present in the
+    # well metadata?
+    well_url, _ = _split_well_path_image_path(new_zarr_url)
+    well_meta = load_NgffWellMeta(well_url)
+    well_paths = [image.path for image in well_meta.well.images]
+
+    if overwrite_input:
+        assert well_paths == ["0"]
+    else:
+        assert well_paths == ["0", "0" + suffix]