changed after to dark before

mantidimaging/core/operations/flat_fielding/flat_fielding.py

@@ -88,33 +88,66 @@ def filter_func(images: ImageStack,
         """
         h.check_data_stack(images)
 
+        # Ensure valid flat-fielding method is selected
+        if selected_flat_fielding not in ["Both, concatenated", "Only Before", "Only After"]:
+            raise ValueError(f"Invalid flat fielding method: {selected_flat_fielding}")
+
+        # Handle missing dark_after by using dark_before if necessary
+        if dark_after is None:
+            dark_after = dark_before
+
+        # Calculate flat and dark averages based on the selected flat-fielding method
         if selected_flat_fielding == "Both, concatenated":
-            if not all([flat_before, flat_after, dark_before, dark_after]):
-                raise ValueError("Missing stack")
+            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'")
+            if dark_before is None and dark_after is None:
+                raise ValueError("Missing stack: dark_before or dark_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:
-                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:
+                if dark_before is not None and dark_after is not None:
+                    dark_avg = (dark_before.data.mean(axis=0) + dark_after.data.mean(axis=0)) / 2.0
+                elif dark_before is not None:
+                    dark_avg = dark_before.data.mean(axis=0)
+                elif dark_after is not None:
+                    dark_avg = dark_after.data.mean(axis=0)
+                else:
+                    dark_avg = np.zeros_like(flat_avg)
+            else:
+                dark_avg = np.zeros_like(flat_avg)
+
+        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)
-        elif selected_flat_fielding == "Only After" and flat_after is not None and dark_after is not None:
+            else:
+                dark_avg = np.zeros_like(flat_avg)
+
+        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:
-            dark_avg = np.zeros_like(flat_avg)
+            else:
+                dark_avg = np.zeros_like(flat_avg)
 
+        # Validate the shapes and execute background correction
         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}")
 

scripts/operations_tests/operations_tests.py

@@ -181,7 +181,7 @@ def add_flats_and_darks(self, params):
         for frame_type in ['flat_before', 'flat_after', 'dark_before', 'dark_after']:
             if frame_type in params:
                 filename_group = FilenameGroup.from_file(config_manager.load_sample)
-                related_filename_group = filename_group.find_related(getattr(FILE_TYPES, frame_type.upper()))
+                related_filename_group = filename_group.find_related(getattr(FILE_TYPES, params[frame_type]))
                 if related_filename_group:
                     related_filename_group.find_all_files()
                     image_stack = loader.load(filename_group=related_filename_group)