Further CLEM bug fixes (#252)

* Load images as 'np.float32' before normalising to 8-bit (np.uin8)
* Save LIF master metadata one level further down
* Run 'process_lif_subimage' in a ProcessPoolExecutor class, and add logic to determine number of workers and threads to use
* Update CPU resources requested by CLEM LIF service
* Use 'tifffile.imread' instead of 'cv2.imread' to open TIFF files
* Added 'imagecodecs' as an explicit dependency; needed for opening certain TIFF files
* Clip after normalising instead of before
* Updated tests

Author: tieneupin

Helm/values.yaml

@@ -47,9 +47,9 @@ clem_process_raw_lifs:
   replicas: 0
   image: gcr.io/image/path
   command: cryoemservices.service -s CLEMProcessRawLIFs -c /cryoemservices/config/cryoemservices_config.yaml
-  cpuRequest: "4"
-  cpuLimit: "4"
-  memoryLimit: 16Gi
+  cpuRequest: "16"
+  cpuLimit: "16"
+  memoryLimit: 32Gi
   scaleOnQueueLength: true
   queueLengthTrigger: "4"
   minReplicaCount: 0

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "graypy",
     "healpy",
     "icebreaker-em",
+    "imagecodecs", # CLEM workflow
     "ispyb>=11.1.2",
     "mrcfile",
     "numpy",

src/cryoemservices/util/image_processing.py

@@ -17,8 +17,8 @@
 import cv2
 import numpy as np
 import SimpleITK as sitk
+import tifffile as tf
 from readlif.reader import LifFile
-from tifffile import imwrite
 
 # Create logger object to output messages with
 logger = logging.getLogger("cryoemservices.util.image_processing")
@@ -61,7 +61,7 @@ class TIFFImageLoader(ImageLoader):
     tiff_file: Path
 
     def load(self) -> np.ndarray:
-        return np.asarray(cv2.imread(self.tiff_file, flags=cv2.IMREAD_UNCHANGED))
+        return tf.imread(self.tiff_file)
 
 
 @dataclass(frozen=True)
@@ -239,7 +239,7 @@ def load_and_convert_image(
     """
 
     try:
-        arr = image_loader.load()
+        arr = image_loader.load().astype(np.float32)
         if new_shape:
             new_y, new_x = new_shape
             arr = cv2.resize(
@@ -248,9 +248,9 @@ def load_and_convert_image(
                 interpolation=cv2.INTER_AREA,
             )
         scale = 255 / ((vmax - vmin) or 1)  # Downscale to 8-bit
-        np.clip(arr, a_min=vmin, a_max=vmax, out=arr)
-        np.subtract(arr, vmin, out=arr, casting="unsafe")
-        np.multiply(arr, scale, out=arr, casting="unsafe")
+        np.subtract(arr, vmin, out=arr)
+        np.multiply(arr, scale, out=arr)
+        np.clip(arr, a_min=0, a_max=255, out=arr)
         return LoadImageResult(
             data=arr.astype(np.uint8),
             frame_num=frame_num,
@@ -342,21 +342,20 @@ def load_and_resize_tile(
         pos_y = int(round((y0 - py0) / (py1 - py0) * parent_y_pixels))
 
         # Load image and resize
-        img = image_loader.load()
+        img = image_loader.load().astype(np.float32)
         resized = cv2.resize(
             img,
             dsize=(tile_x_pixels, tile_y_pixels),
             interpolation=cv2.INTER_AREA,
         )
         # Normalise to 8-bit
         scale = 255 / ((vmax - vmin) or 1)
-        np.clip(resized, vmin, vmax, out=resized)
-        np.subtract(resized, vmin, out=resized, casting="unsafe")
-        np.multiply(resized, scale, out=resized, casting="unsafe")
-        resized = resized.astype(np.uint8)
+        np.subtract(resized, vmin, out=resized)
+        np.multiply(resized, scale, out=resized)
+        np.clip(resized, a_min=0, a_max=255, out=resized)
 
         return ResizeTileResult(
-            data=resized,
+            data=resized.astype(np.uint8),
             frame_num=frame_num,
             x0=pos_x,
             x1=pos_x + tile_x_pixels,
@@ -498,7 +497,7 @@ def write_stack_to_tiff(
     save_name = save_dir.joinpath(file_name + ".tiff")
 
     # With 'bigtiff=True', they have to be pure Python class instances
-    imwrite(
+    tf.imwrite(
         save_name,
         array,
         bigtiff=use_bigtiff,

src/cryoemservices/wrappers/clem_process_raw_lifs.py

@@ -9,7 +9,7 @@
 import json
 import logging
 import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
 from pathlib import Path
 from typing import Optional
 from xml.etree import ElementTree as ET
@@ -429,11 +429,12 @@ def process_lif_file(
     raw_xml_dir = (
         processed_dir
         / "/".join(file.relative_to(processed_dir.parent).parts[1:-1])
+        / file.stem.replace(" ", "_")
         / "metadata"
     )
     for folder in (processed_dir, raw_xml_dir):
         folder.mkdir(parents=True, exist_ok=True)
-        logger.info("Created processing directory and folder to store raw metadata in")
+    logger.info("Created processing directory and folder to store raw metadata in")
 
     # Load LIF file as a LifFile class
     logger.info(f"Loading {file.name!r}")
@@ -464,17 +465,22 @@ def process_lif_file(
     logger.info(f"Examining subimages in {file.name!r}")
 
     # Iterate across the series in the pool
-    results = [
-        process_lif_subimage(
-            file,
-            i,
-            metadata,
-            processed_dir,
-            series_path,
-            num_procs,
-        )
-        for i, (series_path, metadata) in enumerate(metadata_dict.items())
-    ]
+    num_workers = 4 if num_procs > 4 else num_procs
+    num_threads = (num_procs // num_workers) or 1
+    with ProcessPoolExecutor(num_workers) as pool:
+        futures = [
+            pool.submit(
+                process_lif_subimage,
+                file,
+                i,
+                metadata,
+                processed_dir,
+                series_path,
+                num_threads,
+            )
+            for i, (series_path, metadata) in enumerate(metadata_dict.items())
+        ]
+    results = [future.result() for future in futures]
 
     end_time = time.perf_counter()
     logger.debug(f"Processed LIF file {file} in {end_time - start_time}s")
@@ -493,7 +499,7 @@ class ProcessRawLIFsParameters(BaseModel):
 
     lif_file: Path
     root_folder: str  # The root folder under which all LIF files are saved
-    num_procs: int = 20  # Number of processing threads to run
+    num_procs: int = 16  # Number of processing threads to run
 
 
 class ProcessRawLIFsWrapper:

tests/services/test_clem_process_raw_lifs_service.py

@@ -123,7 +123,7 @@ def test_lif_to_stack_service(
     mock_convert.assert_called_with(
         file=lif_file,
         root_folder=raw_dir.stem,
-        number_of_processes=20,
+        number_of_processes=16,
     )
     for result in dummy_results:
         for color in cast(dict[str, Any], result["output_files"]).keys():
@@ -222,7 +222,7 @@ def test_lif_to_stack_service_validation_failed(
     }
     lif_file_value: Any = str(lif_file) if valid_lif_file else 123456789
     root_folder: Any = raw_dir.stem if valid_root_folder else 123456789
-    num_procs: Any = 20 if valid_num_procs else "This is a string"
+    num_procs: Any = 16 if valid_num_procs else "This is a string"
     lif_to_stack_test_message = {
         "lif_file": lif_file_value,
         "root_folder": root_folder,

tests/wrappers/test_clem_process_raw_lifs_wrapper.py

@@ -313,19 +313,24 @@ def test_process_lif_file(
         return_value=raw_xml_metadata,
     )
 
-    # Mock out the sub-image processing function to return results iteratively
-    mock_process_lif_subimage = mocker.patch(
-        "cryoemservices.wrappers.clem_process_raw_lifs.process_lif_subimage"
-    )
-    mock_process_lif_subimage.side_effect = [
-        create_dummy_result(
+    # Store dummy result in mock Future objects
+    mock_futures = []
+    for scene_num in range(num_scenes):
+        mock_future = MagicMock()
+        mock_future.result.return_value = create_dummy_result(
             lif_file=lif_file,
             series_name=series_name,
             processed_dir=processed_dir,
             scene_num=scene_num,
         )
-        for scene_num in range(num_scenes)
-    ]
+        mock_futures.append(mock_future)
+
+    # Mock out the ProcessPoolExecutor used to run 'process_lif_subimage'
+    mock_executor = mocker.patch(
+        "cryoemservices.wrappers.clem_process_raw_lifs.ProcessPoolExecutor"
+    )
+    mock_pool = mock_executor.return_value.__enter__.return_value
+    mock_pool.submit.side_effect = mock_futures
 
     # Run the function
     results = process_lif_file(
@@ -335,7 +340,7 @@ def test_process_lif_file(
     )
 
     # Check that nested list of results was collapsed correctly
-    assert mock_process_lif_subimage.call_count == num_scenes
+    assert mock_pool.submit.call_count == num_scenes
     assert len(results) == num_scenes
 
 
@@ -361,7 +366,7 @@ def test_lif_to_stack_wrapper(
     mock_send_to = mocker.patch("workflows.recipe.wrapper.RecipeWrapper.send_to")
 
     # Set the number of simultaneous processes to run
-    num_procs = 20
+    num_procs = 16
 
     # Construct a dictionary to pass to the wrapper
     message = {