Enhanced Image Stack Loading and Operation Handling for Flat-fielding Tests

mantidimaging/core/operations/flat_fielding/flat_fielding.py

@@ -88,32 +88,50 @@ def filter_func(images: ImageStack,
         """
         h.check_data_stack(images)
 
-        if selected_flat_fielding == "Both, concatenated" and flat_after is not None and flat_before is not None \
-                and dark_after is not None and dark_before is not None:
+        if selected_flat_fielding not in ["Both, concatenated", "Only Before", "Only After"]:
+            raise ValueError(f"Invalid flat fielding method: {selected_flat_fielding}")
+
+        dark_avg = None
+
+        if selected_flat_fielding == "Both, concatenated":
+            if flat_before is None:
+                raise ValueError("Missing stack: flat_before is required for 'Both, concatenated'")
+            if flat_after is None:
+                raise ValueError("Missing stack: flat_after is required for 'Both, concatenated'")
             flat_avg = (flat_before.data.mean(axis=0) + flat_after.data.mean(axis=0)) / 2.0
             if use_dark:
+                if dark_before is None or dark_after is None:
+                    raise ValueError("Missing stack: dark_before and dark_after are required for 'Both, concatenated'")
                 dark_avg = (dark_before.data.mean(axis=0) + dark_after.data.mean(axis=0)) / 2.0
-        elif selected_flat_fielding == "Only Before" and flat_before is not None and dark_before is not None:
-            flat_avg = flat_before.data.mean(axis=0)
-            if use_dark:
-                dark_avg = dark_before.data.mean(axis=0)
-        elif selected_flat_fielding == "Only After" and flat_after is not None and dark_after is not None:
+
+        elif selected_flat_fielding == "Only After":
+            if flat_after is None:
+                raise ValueError("Missing stack: flat_after is required for 'Only After'")
             flat_avg = flat_after.data.mean(axis=0)
             if use_dark:
+                if dark_after is None:
+                    raise ValueError("Missing stack: dark_after is required for 'Only After'")
                 dark_avg = dark_after.data.mean(axis=0)
-        else:
-            raise ValueError("selected_flat_fielding not in:", valid_methods)
 
-        if not use_dark:
+        elif selected_flat_fielding == "Only Before":
+            if flat_before is None:
+                raise ValueError("Missing stack: flat_before is required for 'Only Before'")
+            flat_avg = flat_before.data.mean(axis=0)
+            if use_dark:
+                if dark_before is None:
+                    raise ValueError("Missing stack: dark_before is required for 'Only Before'")
+                dark_avg = dark_before.data.mean(axis=0)
+
+        if dark_avg is None:
             dark_avg = np.zeros_like(flat_avg)
 
         if flat_avg is not None and dark_avg is not None:
-            if 2 != flat_avg.ndim or 2 != dark_avg.ndim:
+            if flat_avg.ndim != 2 or dark_avg.ndim != 2:
                 raise ValueError(
-                    f"Incorrect shape of the flat image ({flat_avg.shape}) or dark image ({dark_avg.shape}) \
-                    which should match the shape of the sample images ({images.data.shape})")
+                    f"Incorrect shape of the flat image ({flat_avg.shape}) or dark image ({dark_avg.shape}) "
+                    f"which should match the shape of the sample images ({images.data.shape[1:]})")
 
-            if not images.data.shape[1:] == flat_avg.shape == dark_avg.shape:
+            if not (images.data.shape[1:] == flat_avg.shape == dark_avg.shape):
                 raise ValueError(f"Not all images are the expected shape: {images.data.shape[1:]}, instead "
                                  f"flat had shape: {flat_avg.shape}, and dark had shape: {dark_avg.shape}")
 
@@ -275,22 +293,6 @@ def _norm_divide(flat: np.ndarray, dark: np.ndarray) -> np.ndarray:
 
 
 def _execute(images: ImageStack, flat=None, dark=None, progress=None):
-    """A benchmark justifying the current implementation, performed on
-    500x2048x2048 images.
-
-    #1 Separate runs
-    Subtract (sequential with np.subtract(data, dark, out=data)) - 13s
-    Divide (par) - 1.15s
-
-    #2 Separate parallel runs
-    Subtract (par) - 5.5s
-    Divide (par) - 1.15s
-
-    #3 Added subtract into _divide so that it is:
-                np.true_divide(
-                    np.subtract(data, dark, out=data), norm_divide, out=data)
-    Subtract then divide (par) - 55s
-    """
     with progress:
         progress.update(msg="Applying background correction")
 

scripts/operations_tests/operations_tests.py

@@ -33,6 +33,7 @@
 from mantidimaging.core.io.loader import loader  # noqa: E402
 from mantidimaging.core.operations.loader import load_filter_packages  # noqa: E402
 from mantidimaging.core.io.instrument_log import InstrumentLog
+from mantidimaging.core.utility.data_containers import FILE_TYPES
 
 script_dir = Path(__file__).resolve().parent
 log_directory = script_dir / "logs"
@@ -144,6 +145,8 @@ def run_test(self, test_case):
         # Handling various pre-run steps
         if test_case.pre_run_step == 'add_nan':
             image_stack = self.add_nan(image_stack, fraction=0.1)
+        elif test_case.pre_run_step == 'add_flats_and_darks':
+            self.add_flats_and_darks(test_case.params)
         elif test_case.pre_run_step == 'load_monitor_log':
             log_data = self.load_monitor_log()
             image_stack.log_file = log_data
@@ -169,6 +172,23 @@ def run_test(self, test_case):
 
         TEST_CASE_RESULTS.append(test_case)
 
+    def add_flats_and_darks(self, params):
+        for frame_type in ['flat_before', 'flat_after', 'dark_before', 'dark_after']:
+            if frame_type not in params:
+                continue
+            try:
+                filename_group = FilenameGroup.from_file(config_manager.load_sample)
+                related_filename_group = filename_group.find_related(getattr(FILE_TYPES, params[frame_type]))
+                if not related_filename_group:
+                    raise ValueError(f"Related files not found for {frame_type}")
+                related_filename_group.find_all_files()
+                image_stack = loader.load(filename_group=related_filename_group)  # Corrected this line
+                if image_stack is None or not hasattr(image_stack, 'data'):
+                    raise ValueError(f"Failed to load or invalid image stack for {frame_type}")
+                params[frame_type] = image_stack
+            except Exception as e:
+                print(f"Error with {frame_type}: {e}")
+
     def load_monitor_log(self):
         filename_group = FilenameGroup.from_file(config_manager.load_sample)
         filename_group.find_log_file()

scripts/operations_tests/test_cases.json

@@ -690,5 +690,56 @@
         }
       }
     ]
+  },
+  "Flat-fielding": {
+    "params": {
+      "use_dark": true,
+      "selected_flat_fielding": "Both, concatenated"
+    },
+    "source_data": "flower128",
+    "cases": [
+      {
+        "test_name": "default_settings",
+        "pre_run_step": "add_flats_and_darks",
+        "params": {
+          "use_dark": true,
+          "selected_flat_fielding": "Both, concatenated",
+          "flat_before": "FLAT_BEFORE",
+          "dark_before": "DARK_BEFORE",
+          "flat_after": "FLAT_AFTER",
+          "dark_after": "DARK_BEFORE"
+        }
+      },
+      {
+        "test_name": "use_only_before",
+        "pre_run_step": "add_flats_and_darks",
+        "params": {
+          "use_dark": true,
+          "selected_flat_fielding": "Only Before",
+          "flat_before": "FLAT_BEFORE",
+          "dark_before": "DARK_BEFORE"
+        }
+      },
+      {
+        "test_name": "use_only_after",
+        "pre_run_step": "add_flats_and_darks",
+        "params": {
+          "use_dark": true,
+          "selected_flat_fielding": "Only After",
+          "flat_after": "FLAT_AFTER",
+          "dark_after": "DARK_BEFORE"
+        }
+      },
+      {
+        "test_name": "no_dark_frames",
+        "pre_run_step": "add_flats_and_darks",
+        "params": {
+          "use_dark": false,
+          "selected_flat_fielding": "Both, concatenated",
+          "flat_before": "FLAT_BEFORE",
+          "flat_after": "FLAT_AFTER"
+        }
+      }
+    ]
   }
 }