compute operation Normalising (#2126)

mantidimaging/core/operations/roi_normalisation/roi_normalisation.py

@@ -3,16 +3,12 @@
 from __future__ import annotations
 
 from functools import partial
-from logging import getLogger
 from typing import TYPE_CHECKING
 
 import numpy as np
-
 from mantidimaging import helper as h
 from mantidimaging.core.operations.base_filter import BaseFilter, FilterGroup
 from mantidimaging.core.parallel import shared as ps
-from mantidimaging.core.parallel import utility as pu
-from mantidimaging.core.utility.progress_reporting import Progress
 from mantidimaging.core.utility.sensible_roi import SensibleROI
 from mantidimaging.gui.utility import add_property_to_form
 from mantidimaging.gui.utility.qt_helpers import Type
@@ -33,11 +29,8 @@ class RoiNormalisationFilter(BaseFilter):
     """Normalises the image data by using the average value in a no-sample (air) region of interest (ROI) of the
     image. This scaling operation allows to account for beam fluctuations and different exposure times of
     projections.
-
     Intended to be used on: Projections
-
     When: Always, to ensure that any fluctuations in beam intensity are normalised.
-
     Note: Beam fluctuations are visible by horizontal lines in the sinograms, as  some projections are
     brighter/darker than the neighbours. This can be fixed with this operation.
     """
@@ -51,44 +44,67 @@ def filter_func(images: ImageStack,
                     flat_field: ImageStack | None = None,
                     progress=None):
         """Normalise by beam intensity.
-
         This does NOT do any checks if the Air Region is out of bounds!
         If the Air Region is out of bounds, the crop will fail at runtime.
         If the Air Region is in bounds, but has overlapping coordinates
         the crop give back a 0 shape of the coordinates that were wrong.
-
         :param images: Sample data which is to be processed. Expected in radiograms
-
         :param region_of_interest: The order is - Left Top Right Bottom. The air
         region for which grey values are summed up and used for normalisation/scaling.
-
         :param normalisation_mode: Controls what the ROI counts are normalised to.
             'Stack Average' : The mean value of the air region across all projections is preserved.
             'Flat Field' : The mean value of the air regions in the projections is made equal to the mean value of the
                            air region in the flat field image.
-
         :param flat_field: Flat field to use if 'Flat Field' mode is enabled.
-
         :param progress: Reference to a progress bar object
-
         :returns: Filtered data (stack of images)
         """
-        if normalisation_mode not in modes():
-            raise ValueError(f"Unknown normalisation_mode: {normalisation_mode}, should be one of {modes()}")
+        if not region_of_interest:
+            raise ValueError('region_of_interest must be provided')
 
-        if normalisation_mode == "Flat Field" and flat_field is None:
-            raise ValueError('flat_field must provided if using normalisation_mode of "Flat Field"')
+        if isinstance(region_of_interest, list):
+            region_of_interest = SensibleROI.from_list(region_of_interest)
 
         h.check_data_stack(images)
 
-        if not region_of_interest:
-            raise ValueError('region_of_interest must be provided')
+        global_params = RoiNormalisationFilter.calculate_global(images, region_of_interest, normalisation_mode,
+                                                                flat_field)
+
+        params = {
+            'region_of_interest': region_of_interest,
+            'normalisation_mode': normalisation_mode,
+            'normalisation_factors': global_params
+        }
+
+        ps.run_compute_func(RoiNormalisationFilter.compute_function, images.data.shape[0], images.shared_array, params)
 
-        progress = Progress.ensure_instance(progress, task_name='ROI Normalisation')
-        _execute(images, region_of_interest, normalisation_mode, flat_field, progress)
         h.check_data_stack(images)
+
         return images
 
+    @staticmethod
+    def calculate_global(images, region_of_interest, normalisation_mode, flat_field):
+        air_means = np.mean(images.data[:, region_of_interest.top:region_of_interest.bottom,
+                                        region_of_interest.left:region_of_interest.right],
+                            axis=(1, 2))
+
+        if normalisation_mode == 'Stack Average':
+            normed_air_means = air_means / air_means.mean()
+        elif normalisation_mode == 'Flat Field' and flat_field is not None:
+            flat_means = np.mean(flat_field.data[:, region_of_interest.top:region_of_interest.bottom,
+                                                 region_of_interest.left:region_of_interest.right],
+                                 axis=(1, 2))
+            normed_air_means = air_means / flat_means.mean()
+        else:
+            raise ValueError(f"Unknown normalisation_mode: {normalisation_mode}")
+
+        return normed_air_means
+
+    @staticmethod
+    def compute_function(i: int, array: np.ndarray, params):
+        normalization_factor = params['normalisation_factors'][i]
+        array[i] /= normalization_factor
+
     @staticmethod
     def register_gui(form, on_change, view):
         label, roi_field = add_property_to_form("Air Region",
@@ -143,63 +159,6 @@ def group_name() -> FilterGroup:
         return FilterGroup.Basic
 
 
-def _calc_mean(data, air_left=None, air_top=None, air_right=None, air_bottom=None):
-    return data[air_top:air_bottom, air_left:air_right].mean()
-
-
-def _divide_by_air(data=None, air_sums=None):
-    data[:] = np.true_divide(data, air_sums)
-
-
-def _execute(images: ImageStack,
-             air_region: SensibleROI,
-             normalisation_mode: str,
-             flat_field: ImageStack | None,
-             progress=None):
-    log = getLogger(__name__)
-
-    with progress:
-        progress.update(msg="Normalization by air region")
-        if isinstance(air_region, list):
-            air_region = SensibleROI.from_list(air_region)
-
-        # initialise same number of air sums
-        img_num = images.data.shape[0]
-        air_means = pu.create_array((img_num, ), images.dtype)
-
-        do_calculate_air_means = ps.create_partial(_calc_mean,
-                                                   ps.return_to_second_at_i,
-                                                   air_left=air_region.left,
-                                                   air_top=air_region.top,
-                                                   air_right=air_region.right,
-                                                   air_bottom=air_region.bottom)
-
-        arrays = [images.shared_array, air_means]
-        ps.execute(do_calculate_air_means, arrays, images.data.shape[0], progress)
-
-        if normalisation_mode == 'Stack Average':
-            air_means.array /= air_means.array.mean()
-
-        elif normalisation_mode == 'Flat Field' and flat_field is not None:
-            flat_mean = pu.create_array((flat_field.data.shape[0], ), flat_field.dtype)
-            arrays = [flat_field.shared_array, flat_mean]
-            ps.execute(do_calculate_air_means, arrays, flat_field.data.shape[0], progress)
-            air_means.array /= flat_mean.array.mean()
-
-        if np.isnan(air_means.array).any():
-            raise ValueError("Air region contains invalid (NaN) pixels")
-
-        do_divide = ps.create_partial(_divide_by_air, fwd_function=ps.inplace2)
-        arrays = [images.shared_array, air_means]
-        ps.execute(do_divide, arrays, images.data.shape[0], progress)
-
-        avg = np.average(air_means.array)
-        max_avg = np.max(air_means.array) / avg
-        min_avg = np.min(air_means.array) / avg
-
-        log.info(f"Normalization by air region. Average: {avg}, max ratio: {max_avg}, min ratio: {min_avg}.")
-
-
 def enable_correct_fields_only(text, flat_file_widget):
     if text == "Flat Field":
         flat_file_widget.setEnabled(True)