Add apply_masks option to single-cell image extraction

src/scportrait/pipeline/extraction.py

@@ -68,6 +68,9 @@ def __init__(self, *args, **kwargs):
         self.extraction_file = os.path.join(self.directory, self.DEFAULT_DATA_DIR, self.DEFAULT_EXTRACTION_FILE)
         self.output_path = None
 
+        # whether to mask out background/neighbouring cells when extracting image channels
+        self.apply_masks = True
+
     def _get_compression_type(self) -> None:
         """setup compression of single-cell images in HDF5 file based on config."""
         # default value for compression is "lzf" is nothing else is specified
@@ -640,9 +643,11 @@ def _extract_classes(
         else:
             image_data = self.image_data[image_index, :, window_y, window_x]
 
-        image_data = (
-            image_data * masks[-1]
-        )  # always uses the last available mask, in nucleus only seg its the nucleus, if both its the cytosol, if only cytosol its also the cytosol. This always is the mask we want to use to extract the channel information
+        if self.apply_masks:
+            image_data = (
+                image_data * masks[-1]
+            )  # uses the last available mask, in nucleus only seg its the nucleus, if both its the cytosol, if only cytosol its also the cytosol. This always is the mask we want to use to extract the channel information
+        # if apply_masks is False the raw image window is kept unmasked (background/neighbouring cells retained)
 
         # this needs to be performed on a per channel basis!
         images = []
@@ -1143,7 +1148,12 @@ def _calibrate_max_inflight_result_batches(
         return calibrated_max, first_wait_s, first_result, pending_results, next_submit_ix
 
     def process(
-        self, partial: bool = False, n_cells: int = None, seed: int = 42, output_folder_name: str | None = None
+        self,
+        partial: bool = False,
+        n_cells: int = None,
+        seed: int = 42,
+        output_folder_name: str | None = None,
+        apply_masks: bool = True,
     ) -> None:
         """
         Extracts single cell images from a segmented scPortrait project and saves the results to a standardized HDF5 file.
@@ -1153,6 +1163,9 @@ def process(
             partial: if set to True only a random subset of n_cells will be extracted.
             n_cells: Number of cells to extract if partial is set to True.
             seed: Seed for random sampling of cells for reproducibility if partial is set to True.
+            apply_masks: if set to True (default) the segmentation mask is applied to the extracted
+                image channels so that background and neighbouring cells are zeroed out. If set to
+                False the raw image window is saved without masking.
 
         Important
         ---------
@@ -1240,6 +1253,9 @@ def process(
 
         start_setup = timeit.default_timer()
 
+        # set up flag for whether to mask image channels (set before pool creation so it propagates to workers)
+        self.apply_masks = apply_masks
+
         # set up flag for partial processing
         self.partial_processing = partial
         if self.partial_processing:

src/scportrait/pipeline/project.py

@@ -1646,6 +1646,7 @@ def extract(
         seed: int = 42,
         overwrite: bool | None = None,
         output_folder_name: str | None = None,
+        apply_masks: bool = True,
     ) -> None:
         """Extract single-cell images from the input image using the defined extraction method.
 
@@ -1656,6 +1657,9 @@ def extract(
             overwrite: If set to ``None``, will read the overwrite value from the associated project.
                 Otherwise can be set to a boolean value to override project specific settings for image loading
             output_folder_name: can be set to override the default output folder location.
+            apply_masks: If set to ``True`` (default), the segmentation mask is applied to the extracted
+                image channels so background and neighbouring cells are zeroed out (current behaviour).
+                If set to ``False``, the raw image window is extracted without masking.
 
         Returns:
             Single-cell images are written to HDF5 files in the project associated extraction directory. The file path can be accessed via ``project.extraction_f.output_path``.
@@ -1674,7 +1678,13 @@ def extract(
         if overwrite is not None:
             self.extraction_f.overwrite_run_path = overwrite
 
-        self.extraction_f(partial=partial, n_cells=n_cells, seed=seed, output_folder_name=output_folder_name)
+        self.extraction_f(
+            partial=partial,
+            n_cells=n_cells,
+            seed=seed,
+            output_folder_name=output_folder_name,
+            apply_masks=apply_masks,
+        )
         self.get_project_status()
 
     def featurize(