From a6bd16bfe87d629c1afd57d0aa6e441e4bf554a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ali=20O=C4=9Fuz=20Can?= Date: Fri, 3 Jul 2026 17:52:05 +0200 Subject: [PATCH] Add apply_masks option to single-cell image extraction HDF5CellExtraction previously always multiplied the extracted image channels by the segmentation mask, zeroing out background and neighbouring cells. This adds an `apply_masks` parameter (default True, preserving current behaviour) to `Project.extract` and `HDF5CellExtraction.process`. When set to False, the raw image window is saved without masking, enabling downstream analysis of unmasked single-cell images. The mask channel(s) are still stored either way. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/scportrait/pipeline/extraction.py | 24 ++++++++++++++++++++---- src/scportrait/pipeline/project.py | 12 +++++++++++- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/scportrait/pipeline/extraction.py b/src/scportrait/pipeline/extraction.py index 32b2eca5..7c6a05f4 100644 --- a/src/scportrait/pipeline/extraction.py +++ b/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: diff --git a/src/scportrait/pipeline/project.py b/src/scportrait/pipeline/project.py index 74433e06..e6036135 100644 --- a/src/scportrait/pipeline/project.py +++ b/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(