Merge pull request #6 from ornlneutronimaging/re_filter_fix

Address re-filtering sigma estimate issue

.github/workflows/unittest.yml

@@ -28,7 +28,7 @@ jobs:
           echo "installing additional dependencies if cannot be installed from conda"
       - name: run unit tests
         run: |
-          echo "running unit tests"
+          echo "running unit tests (CPU)"
           python -m pytest --cov=src --cov-report=xml --cov-report=term-missing -m "not cuda_required"
       - name: upload coverage to codecov
         uses: codecov/codecov-action@v4
@@ -44,3 +44,27 @@ jobs:
           CHANNELS="--channel mantid/label/main --channel conda-forge"
           VERSION=$(versioningit ../) conda mambabuild $CHANNELS --output-folder . .
           conda verify noarch/bm3dornl*.tar.bz2
+
+  linux_gpu:
+    runs-on: github-gpu-builder
+    defaults:
+      run:
+        shell: bash -l {0}
+    steps:
+      - uses: actions/checkout@v3
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          miniforge-version: "latest"
+          auto-update-conda: true
+          channels: conda-forge,defaults
+          mamba-version: "*"
+          environment-file: environment.yml
+          cache-environment-key: ${{ runner.os }}-env-${{ hashFiles('**/environment.yml') }}
+          cache-downloads-key: ${{ runner.os }}-downloads-${{ hashFiles('**/environment.yml') }}
+      - name: install additional dependencies
+        run: |
+          echo "installing additional dependencies if cannot be installed from conda"
+      - name: run unit tests
+        run: |
+          echo "running unit tests (GPU)"
+          python -m pytest --cov=src --cov-report=xml --cov-report=term-missing -m "cuda_required"

src/bm3dornl/denoiser.py

@@ -15,6 +15,7 @@
 from bm3dornl.utils import (
     horizontal_binning,
     horizontal_debinning,
+    estimate_noise_std,
 )
 
 
@@ -160,8 +161,9 @@ def re_filtering(
         weights = np.zeros_like(self.image, dtype=np.float64)
 
         # estimate the noise
-        noise = np.asarray(self.image) - self.estimate_denoised_image
-        sigma_squared = np.mean(noise**2)
+        sigma_squared = estimate_noise_std(
+            noisy_image=self.image, noise_free_image=self.estimate_denoised_image
+        )
 
         logging.info("Block matching for 2nd pass...")
         block, positions = self.patch_manager.get_hyper_block(
@@ -171,7 +173,7 @@ def re_filtering(
         )
 
         logging.info("Wiener-Hadamard filtering...")
-        block = wiener_hadamard(block, sigma_squared * 1e3)  # why does this work?
+        block = wiener_hadamard(block, sigma_squared)  # why does this work?
 
         # manual release of memory
         memory_cleanup()
@@ -214,12 +216,12 @@ def denoise(
         self.thresholding(
             cut_off_distance, intensity_diff_threshold, num_patches_per_group, threshold
         )
-        self.final_denoised_image = self.estimate_denoised_image
+        # self.final_denoised_image = self.estimate_denoised_image
 
-        # logging.info("Second pass: Re-filtering")
-        # self.re_filtering(
-        #     cut_off_distance, intensity_diff_threshold, num_patches_per_group
-        # )
+        logging.info("Second pass: Re-filtering")
+        self.re_filtering(
+            cut_off_distance, intensity_diff_threshold, num_patches_per_group
+        )
 
 
 def bm3d_streak_removal(

src/bm3dornl/utils.py

@@ -267,3 +267,42 @@ def horizontal_debinning(original_image: np.ndarray, target: np.ndarray) -> np.n
     interpolated_image = spline(new_y, new_x)
 
     return interpolated_image
+
+
+def estimate_noise_std(
+    noisy_image: np.ndarray,
+    noise_free_image: np.ndarray,
+) -> float:
+    """
+    Estimate the noise standard deviation from a pair of noisy and noise-free images.
+
+    Parameters
+    ----------
+    noisy_image : np.ndarray
+        A 2D NumPy array representing the noisy image.
+    noise_free_image : np.ndarray
+        A 2D NumPy array representing the noise-free image.
+
+    Returns
+    -------
+    float
+        The estimated noise standard deviation.
+    """
+    # Rescale both images to [0, 255]
+    noisy_image_min, noisy_image_max = noisy_image.min(), noisy_image.max()
+    noise_free_image_min, noise_free_image_max = (
+        noise_free_image.min(),
+        noise_free_image.max(),
+    )
+
+    noisy_image = (
+        (noisy_image - noisy_image_min) / (noisy_image_max - noisy_image_min) * 255
+    )
+    noise_free_image = (
+        (noise_free_image - noise_free_image_min)
+        / (noise_free_image_max - noise_free_image_min)
+        * 255
+    )
+
+    # calculate the noise standard deviation
+    return np.std(np.abs(noisy_image - noise_free_image))

tests/unit/bm3dornl/test_utils.py

@@ -5,6 +5,7 @@
 import pytest
 import numpy as np
 from bm3dornl.utils import (
+    estimate_noise_std,
     find_candidate_patch_ids,
     is_within_threshold,
     get_signal_patch_positions,
@@ -229,5 +230,24 @@ def test_horizontal_debinning_scaling(original_width, target_width):
     ), f"Failed to scale from {original_width} to {target_width}"
 
 
+def test_estimate_noise_std():
+    # Create a random noise-free image
+    noise_free_image = np.random.rand(8, 8)
+
+    # Add Gaussian noise to the noise-free image
+    noise_std_true = 0.2
+    noisy_image = noise_free_image + np.random.normal(
+        0.5, noise_std_true, noise_free_image.shape
+    )
+
+    # Estimate the noise standard deviation using the function
+    estimated_noise_std = estimate_noise_std(noisy_image, noise_free_image) / 255
+
+    # Assert that the estimate is close to the true value, with a high tolerance
+    assert np.isclose(
+        estimated_noise_std, noise_std_true, atol=0.2
+    ), "Estimated noise standard deviation is incorrect"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])