diff --git a/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md
new file mode 100644
index 00000000..82b6f9c6
--- /dev/null
+++ b/changelog.d/20260909_155518_chaitanyachawak001_225_separate_noise_estimation_from_sample_weight_computation.md
@@ -0,0 +1,15 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+For top level release notes, leave all the headers commented out.
+-->
+
+### New features
+
+- Added `NoiseEstimator.estimate_noise_batch` for batch noise estimation, which is reusable outside training (e.g. SHE validation) since it no longer depends on the training loss/output representation.
+
+### Internal changes
+
+- Moved `NoiseEstimator` from `utils.utils` to a dedicated `utils.noise` module, and initialised it with a default exclusion-window radius (`NoiseEstimator.default_win_rad`).
+- Simplified `train_utils.calculate_sample_weights` to `calculate_sample_weights(images, masks=None, ...)`. It no longer needs to know about `loss` or `use_sample_weights`. `general_train_cycle` now resolves the training output representation based on `loss.name`: any name starting with the `masked_` prefix (e.g. `masked_mean_squared_error`) is treated as a masked loss, so new masked loss functions must follow this naming convention to be handled correctly.
diff --git a/src/wf_psf/tests/test_training/train_utils_test.py b/src/wf_psf/tests/test_training/train_utils_test.py
index db05d8f3..b412c946 100644
--- a/src/wf_psf/tests/test_training/train_utils_test.py
+++ b/src/wf_psf/tests/test_training/train_utils_test.py
@@ -31,9 +31,9 @@ def mock_noise_estimator():
     """Mock the NoiseEstimator class to return a fixed noise estimate."""
     with patch("wf_psf.training.train_utils.NoiseEstimator") as MockNoiseEstimator:
         mock_instance = MockNoiseEstimator.return_value
-        mock_instance.estimate_noise.side_effect = lambda img, mask=None: np.std(
-            img
-        )  # Mock behavior
+        mock_instance.estimate_noise_batch.side_effect = (
+            lambda images, masks=None: np.array([np.std(img) for img in images])
+        )
         yield mock_instance
 
 
@@ -118,72 +118,31 @@ def test_learning_rate_affects_default_optimizer():
 
 
 @pytest.mark.parametrize(
-    "use_sample_weights, expected_output",
-    [
-        (False, None),
-        (True, np.ndarray),  # When enabled, it should return an array
-    ],
-)
-def test_calculate_sample_weights_no_weights(
-    mock_noise_estimator, use_sample_weights, expected_output
-):
-    """Test when sample weights are disabled and ensure correct return type."""
-    outputs = np.random.rand(5, 32, 32)  # (batch_size=5, height=32, width=32)
-    result = train_utils.calculate_sample_weights(
-        outputs, use_sample_weights, loss=None
-    )
-    if not use_sample_weights:
-        assert result is None
-    else:
-        assert isinstance(result, expected_output)
-
-
-@pytest.mark.parametrize(
-    "use_sample_weights, loss, expected_output_type",
+    "masks",
     [
-        (False, None, None),  # Null test: when disabled, should return None
-        (True, "mean_squared_error", np.ndarray),  # Standard MSE case
-        (True, "masked_mean_squared_error", np.ndarray),  # Masked MSE case
+        None,
+        np.random.randint(0, 2, size=(5, 32, 32)).astype(bool),  # Per-image mask
     ],
 )
-def test_calculate_sample_weights_integration(
-    use_sample_weights, loss, expected_output_type
-):
-    """Test different cases for sample weight computation."""
-    # Generate dummy image data
+def test_calculate_sample_weights_integration(masks):
+    """Test sample weight computation with real (unmocked) noise estimation."""
     batch_size, height, width = 5, 32, 32
+    images = np.random.rand(batch_size, height, width)
 
-    if loss == "masked_mean_squared_error":
-        # Create image-mask pairs: last dimension has [image, mask]
-        outputs = np.random.rand(batch_size, height, width, 2)
-        outputs[..., 1] = np.random.randint(
-            0, 2, size=(batch_size, height, width)
-        )  # Binary mask
-    else:
-        outputs = np.random.rand(batch_size, height, width)
-
-    result = train_utils.calculate_sample_weights(outputs, use_sample_weights, loss)
+    result = train_utils.calculate_sample_weights(images, masks=masks)
 
-    if not use_sample_weights:
-        assert result is None
-    else:
-        assert isinstance(result, expected_output_type)
-        assert result.shape == (batch_size,)  # Expecting one weight per sample
-        assert np.all(result > 0)  # Weights should be positive
-        assert np.isclose(
-            np.median(result), 1, atol=0.1
-        )  # Weights should be scaled around 1
+    assert isinstance(result, np.ndarray)
+    assert result.shape == (batch_size,)  # Expecting one weight per sample
+    assert np.all(result > 0)  # Weights should be positive
+    assert np.isclose(
+        np.median(result), 1, atol=0.1
+    )  # Weights should be scaled around 1
 
 
-@pytest.mark.parametrize(
-    "loss", [None, "mean_squared_error", "masked_mean_squared_error"]
-)
-def test_calculate_sample_weights_unit(mock_noise_estimator, loss):
+def test_calculate_sample_weights_unit(mock_noise_estimator):
     """Test sample weighting strategy with random images."""
-    outputs = np.random.rand(10, 32, 32)  # 10 images of size 32x32
-    result = train_utils.calculate_sample_weights(
-        outputs, use_sample_weights=True, loss=loss
-    )
+    images = np.random.rand(10, 32, 32)  # 10 images of size 32x32
+    result = train_utils.calculate_sample_weights(images)
 
     # Check the type and shape of the result
     assert isinstance(result, np.ndarray)
@@ -196,27 +155,19 @@ def test_calculate_sample_weights_unit(mock_noise_estimator, loss):
     assert np.isclose(np.median(result), 1, atol=0.1)
 
 
-@pytest.mark.parametrize(
-    "loss", [None, "mean_squared_error", "masked_mean_squared_error"]
-)
-def test_calculate_sample_weights_high_variance(mock_noise_estimator, loss):
+def test_calculate_sample_weights_high_variance(mock_noise_estimator):
     """Test case for high variance (noisy images)."""
     # Create high variance images with more noise
-    outputs = np.random.normal(loc=0.0, scale=10.0, size=(5, 32, 32))  # Larger noise
+    images = np.random.normal(loc=0.0, scale=10.0, size=(5, 32, 32))  # Larger noise
 
     # Calculate sample weights
-    weights = train_utils.calculate_sample_weights(
-        outputs, use_sample_weights=True, loss=loss
-    )
+    weights = train_utils.calculate_sample_weights(images)
 
-    # Check that weights are relatively lower for high variance images
     # Compare with weights from a small variance case
-    small_variance_outputs = np.random.normal(
+    small_variance_images = np.random.normal(
         loc=0.0, scale=1e-6, size=(5, 32, 32)
     )  # Small noise
-    small_variance_weights = train_utils.calculate_sample_weights(
-        small_variance_outputs, use_sample_weights=True, loss=None
-    )
+    small_variance_weights = train_utils.calculate_sample_weights(small_variance_images)
 
     # Check if the median of the high variance weights is smaller than the median of the low variance weights
     assert np.median(weights) <= np.median(small_variance_weights), (
@@ -228,6 +179,70 @@ def test_calculate_sample_weights_high_variance(mock_noise_estimator, loss):
     assert np.all(weights >= 0), f"Weights are negative: {weights}"
 
 
+@pytest.mark.parametrize(
+    "loss_name",
+    [
+        None,
+        "mean_squared_error",
+        "some_other_loss",
+        "non_masked_loss",
+        "unmasked_mse_loss",
+        "masked",
+    ],
+)
+def test_resolve_training_outputs_unmasked_loss(loss_name):
+    """Test that for a non-"masked_"-prefixed loss name, outputs pass through unchanged and masks is None."""
+    outputs = np.random.rand(4, 8, 8)
+    images, masks = train_utils._resolve_training_outputs(outputs, loss_name)
+
+    assert images is outputs
+    assert masks is None
+
+
+@pytest.mark.parametrize(
+    "loss_name", ["masked_mean_squared_error", "masked_huber_error"]
+)
+def test_resolve_training_outputs_masked_loss(loss_name):
+    """Test that any loss name starting with "masked_" splits outputs into images
+    and an inverted boolean mask, so new masked losses work without code changes."""
+    batch_size, height, width = 4, 8, 8
+    images = np.random.rand(batch_size, height, width)
+    raw_masks = np.random.randint(0, 2, size=(batch_size, height, width))
+    outputs = np.stack([images, raw_masks], axis=-1)
+
+    resolved_images, resolved_masks = train_utils._resolve_training_outputs(
+        outputs, loss_name
+    )
+
+    np.testing.assert_array_equal(resolved_images, images)
+    np.testing.assert_array_equal(resolved_masks, np.array(1 - raw_masks, dtype=bool))
+
+
+def test_calculate_sample_weights_uses_noise_estimator_batch():
+    """Test that calculate_sample_weights delegates noise estimation to NoiseEstimator.estimate_noise_batch."""
+    batch_size, height, width = 5, 32, 32
+    images = np.random.rand(batch_size, height, width)
+    masks = np.random.randint(0, 2, size=(batch_size, height, width)).astype(bool)
+
+    fake_sigma = np.linspace(1.0, 2.0, batch_size)
+    with patch("wf_psf.training.train_utils.NoiseEstimator") as MockNoiseEstimator:
+        mock_instance = MockNoiseEstimator.return_value
+        mock_instance.estimate_noise_batch.return_value = fake_sigma
+
+        result = train_utils.calculate_sample_weights(images, masks)
+
+    mock_instance.estimate_noise_batch.assert_called_once()
+    # Passed the same images and masks through, with no knowledge of any loss
+    called_images, called_masks = mock_instance.estimate_noise_batch.call_args[0][:2]
+    assert called_images is images
+    assert called_masks is masks
+
+    # Weights are the median-normalised inverse variance of the returned sigma
+    expected = 1 / fake_sigma**2
+    expected /= np.median(expected)
+    np.testing.assert_allclose(result, expected)
+
+
 @pytest.mark.parametrize(
     "loss_function, metric_function",
     [
@@ -365,7 +380,76 @@ def test_general_train_cycle_handles_cycle_def(
     if non_param_train:
         expected_calls.append(call(param_bool=False, nonparam_bool=True))
 
-    mock_model.set_trainable_layers.assert_has_calls(expected_calls, any_order=False)
+    mock_model.set_trainable_layers.asserFootert_has_calls(expected_calls, any_order=False)
+
+
+def test_general_train_cycle_resolves_images_and_masks_for_masked_loss(
+    mock_test_setup,
+):
+    """Test that general_train_cycle unpacks image/mask from outputs when the loss is masked,
+    before calculate_sample_weights ever sees the data."""
+    mock_model = mock_test_setup["mock_model"]
+    inputs = mock_test_setup["inputs"]
+
+    batch_size, height, width = 4, 8, 8
+    images = np.random.rand(batch_size, height, width)
+    raw_masks = np.random.randint(0, 2, size=(batch_size, height, width))
+    outputs = np.stack([images, raw_masks], axis=-1)
+
+    with patch(
+        "wf_psf.training.train_utils.calculate_sample_weights"
+    ) as mock_calculate_sample_weights:
+        mock_calculate_sample_weights.return_value = np.ones(batch_size)
+        train_utils.general_train_cycle(
+            psf_model=mock_model,
+            inputs=inputs,
+            outputs=outputs,
+            validation_data=mock_test_setup["validation_data"],
+            batch_size=mock_test_setup["batch_size"],
+            learning_rate_param=mock_test_setup["learning_rate_param"],
+            learning_rate_non_param=mock_test_setup["learning_rate_non_param"],
+            n_epochs_param=mock_test_setup["n_epochs_param"],
+            n_epochs_non_param=mock_test_setup["n_epochs_non_param"],
+            param_loss=train_utils.MaskedMeanSquaredError(),
+            use_sample_weights=True,
+            cycle_def="parametric",
+        )
+
+    mock_calculate_sample_weights.assert_called_once()
+    called_images, called_masks = mock_calculate_sample_weights.call_args[0][:2]
+    np.testing.assert_array_equal(called_images, images)
+    np.testing.assert_array_equal(called_masks, np.array(1 - raw_masks, dtype=bool))
+
+
+def test_general_train_cycle_skips_sample_weights_when_disabled(mock_test_setup):
+    """Test that general_train_cycle leaves sample_weight as None and never calls
+    calculate_sample_weights when use_sample_weights is False."""
+    mock_model = mock_test_setup["mock_model"]
+
+    with (
+        patch(
+            "wf_psf.training.train_utils.calculate_sample_weights"
+        ) as mock_calculate_sample_weights,
+        patch("wf_psf.training.train_utils.train_cycle_part") as mock_train_cycle_part,
+    ):
+        mock_train_cycle_part.return_value = MagicMock()
+        train_utils.general_train_cycle(
+            psf_model=mock_model,
+            inputs=mock_test_setup["inputs"],
+            outputs=mock_test_setup["outputs"],
+            validation_data=mock_test_setup["validation_data"],
+            batch_size=mock_test_setup["batch_size"],
+            learning_rate_param=mock_test_setup["learning_rate_param"],
+            learning_rate_non_param=mock_test_setup["learning_rate_non_param"],
+            n_epochs_param=mock_test_setup["n_epochs_param"],
+            n_epochs_non_param=mock_test_setup["n_epochs_non_param"],
+            use_sample_weights=False,
+            cycle_def="parametric",
+        )
+
+    mock_calculate_sample_weights.assert_not_called()
+    _, train_cycle_part_kwargs = mock_train_cycle_part.call_args
+    assert train_cycle_part_kwargs["sample_weight"] is None
 
 
 def test_get_callbacks():
diff --git a/src/wf_psf/tests/test_utils/noise_test.py b/src/wf_psf/tests/test_utils/noise_test.py
new file mode 100644
index 00000000..6a20302d
--- /dev/null
+++ b/src/wf_psf/tests/test_utils/noise_test.py
@@ -0,0 +1,175 @@
+import numpy as np
+from wf_psf.utils.noise import NoiseEstimator
+
+
+def test_initialization():
+    """Test if NoiseEstimator initializes correctly."""
+    img_dim = (50, 50)
+    win_rad = 10
+    estimator = NoiseEstimator(img_dim, win_rad)
+
+    assert estimator.img_dim == img_dim
+    assert estimator.win_rad == win_rad
+    assert isinstance(estimator.window, np.ndarray)
+    assert estimator.window.shape == img_dim
+
+
+def test_default_win_rad():
+    """Test that the default_win_rad derives an int radius from the image dimensions."""
+    assert NoiseEstimator.default_win_rad((32, 32)) == 10
+    assert NoiseEstimator.default_win_rad((40, 40)) == 13
+    assert NoiseEstimator.default_win_rad((100, 76)) == 31
+    assert isinstance(NoiseEstimator.default_win_rad((32, 32)), int)
+
+
+def test_initialization_uses_default_win_rad_when_omitted():
+    """Test that omitting win_rad falls back to default_win_rad and produces the same window."""
+    img_dim = (50, 50)
+    estimator = NoiseEstimator(img_dim)
+    explicit_estimator = NoiseEstimator(img_dim, NoiseEstimator.default_win_rad(img_dim))
+
+    assert estimator.win_rad == NoiseEstimator.default_win_rad(img_dim)
+    assert isinstance(estimator.win_rad, int)
+    np.testing.assert_array_equal(estimator.window, explicit_estimator.window)
+
+
+def test_init_window():
+    """Test that the exclusion window is correctly initialized."""
+    img_dim = (50, 50)
+    win_rad = 10
+    estimator = NoiseEstimator(img_dim, win_rad)
+
+    mid_x, mid_y = img_dim[0] / 2, img_dim[1] / 2
+
+    for x in range(img_dim[0]):
+        for y in range(img_dim[1]):
+            # Pixel inside the exclusion radius should be False, others True
+            inside_radius = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2) <= win_rad
+            assert estimator.window[x, y] == (not inside_radius)
+
+
+def test_sigma_mad():
+    """Test the MAD-based standard deviation estimation."""
+    data = np.array(
+        [1, 1, 2, 2, 3, 3, 4, 4, 100]
+    )  # Outlier should not heavily influence MAD
+    expected_sigma = 1.4826 * np.median(np.abs(data - np.median(data)))
+
+    assert np.isclose(NoiseEstimator.sigma_mad(data), expected_sigma, atol=1e-4)
+
+
+def test_estimate_noise_without_default_window():
+    """Test noise estimation with the default exclusion window (no custom mask)."""
+    img_dim = (50, 50)
+    win_rad = 5
+    estimator = NoiseEstimator(img_dim, win_rad)
+
+    # Create a synthetic noisy image (Gaussian noise with mean=0, std=10)
+    np.random.seed(42)
+    image = np.random.normal(0, 10, img_dim)
+
+    noise_estimation = estimator.estimate_noise(image)
+
+    # The estimated noise should be close to 10 (the true std)
+    assert np.isclose(noise_estimation, 10, atol=2)
+
+
+def test_estimate_noise_with_custom_mask():
+    """Test noise estimation with a custom mask applied outside the exclusion radius."""
+    img_dim = (50, 50)
+    estimator = NoiseEstimator(img_dim, win_rad=5)
+
+    # Create synthetic noise with std=5
+    np.random.seed(42)
+    image = np.random.normal(0, 5, img_dim)
+
+    # Custom mask excluding top-left corner
+    custom_mask = np.ones(img_dim, dtype=bool)
+    custom_mask[:10, :10] = False  # Mask out top-left 10x10 pixels
+
+    noise_estimation = estimator.estimate_noise(image, mask=custom_mask)
+
+    assert np.isclose(noise_estimation, 5, atol=1)
+
+
+def test_apply_mask_with_none_mask():
+    """Test apply_mask when mask is None."""
+    img_dim = (10, 10)
+    estimator = NoiseEstimator(img_dim, win_rad=3)
+
+    result = estimator.apply_mask(None)  # Pass None as the mask
+
+    # It should return the window itself when no mask is provided
+    assert np.array_equal(
+        result, estimator.window
+    ), "apply_mask should return the window when mask is None."
+
+
+def test_apply_mask_with_valid_mask():
+    """Test apply_mask when a valid mask is provided."""
+    img_dim = (10, 10)
+    estimator = NoiseEstimator(img_dim, win_rad=3)
+
+    # Create a custom mask
+    custom_mask = np.ones(img_dim, dtype=bool)
+    custom_mask[5, 5] = False  # Set a pixel to False to exclude it from the window
+
+    result = estimator.apply_mask(custom_mask)
+
+    # Check that the mask was applied correctly: pixel (5, 5) should be False, others True
+    expected_result = estimator.window & custom_mask
+    assert np.array_equal(
+        result, expected_result
+    ), "apply_mask did not apply the mask correctly."
+
+
+def test_apply_mask_with_zeroed_mask():
+    """Test apply_mask when a zeroed mask is provided."""
+    img_dim = (10, 10)
+    estimator = NoiseEstimator(img_dim, win_rad=3)
+
+    # Create a mask where all pixels are excluded (False)
+    zeroed_mask = np.zeros(img_dim, dtype=bool)
+
+    result = estimator.apply_mask(zeroed_mask)
+
+    # The result should be an array of False values, as the mask excludes all pixels
+    expected_result = np.zeros(img_dim, dtype=bool)
+    assert np.array_equal(
+        result, expected_result
+    ), "apply_mask did not handle the zeroed mask correctly."
+
+
+def test_estimate_noise_batch_without_masks():
+    """estimate_noise_batch with no masks matches per-image estimate_noise."""
+    img_dim = (30, 30)
+    win_rad = 5
+    estimator = NoiseEstimator(img_dim, win_rad)
+
+    np.random.seed(0)
+    images = np.random.normal(0, 3, size=(4, *img_dim))
+
+    batch_result = estimator.estimate_noise_batch(images)
+    expected = np.array([estimator.estimate_noise(img) for img in images])
+
+    assert batch_result.shape == (4,)
+    np.testing.assert_allclose(batch_result, expected)
+
+
+def test_estimate_noise_batch_with_masks():
+    """estimate_noise_batch with per-image masks matches per-image estimate_noise."""
+    img_dim = (30, 30)
+    win_rad = 5
+    estimator = NoiseEstimator(img_dim, win_rad)
+
+    np.random.seed(1)
+    images = np.random.normal(0, 3, size=(4, *img_dim))
+    masks = np.random.randint(0, 2, size=(4, *img_dim)).astype(bool)
+
+    batch_result = estimator.estimate_noise_batch(images, masks)
+    expected = np.array(
+        [estimator.estimate_noise(img, mask) for img, mask in zip(images, masks)]
+    )
+
+    assert batch_result.shape == (4,)
+    np.testing.assert_allclose(batch_result, expected)
diff --git a/src/wf_psf/tests/test_utils/utils_test.py b/src/wf_psf/tests/test_utils/utils_test.py
index cc7f2a2b..c849c692 100644
--- a/src/wf_psf/tests/test_utils/utils_test.py
+++ b/src/wf_psf/tests/test_utils/utils_test.py
@@ -11,7 +11,6 @@
 import numpy as np
 from wf_psf.utils.utils import (
     downsample_im,
-    NoiseEstimator,
     zernike_generator,
     compute_unobscured_zernike_projection,
     decompose_tf_obscured_opd_basis,
@@ -19,143 +18,6 @@
 from wf_psf.sims.psf_simulator import PSFSimulator
 from unittest import mock
 
-
-def test_downsample_basic():
-    """Test apply_mask when a zeroed mask is provided."""
-    img_dim = (10, 10)
-    estimator = NoiseEstimator(img_dim, win_rad=3)
-
-    # Create a mask where all pixels are excluded (False)
-    zeroed_mask = np.zeros(img_dim, dtype=bool)
-
-    result = estimator.apply_mask(zeroed_mask)
-
-    # The result should be an array of False values, as the mask excludes all pixels
-    expected_result = np.zeros(img_dim, dtype=bool)
-    assert np.array_equal(
-        result, expected_result
-    ), "apply_mask did not handle the zeroed mask correctly."
-
-
-def test_initialization():
-    """Test if NoiseEstimator initializes correctly."""
-    img_dim = (50, 50)
-    win_rad = 10
-    estimator = NoiseEstimator(img_dim, win_rad)
-
-    assert estimator.img_dim == img_dim
-    assert estimator.win_rad == win_rad
-    assert isinstance(estimator.window, np.ndarray)
-    assert estimator.window.shape == img_dim
-
-
-def test_init_window():
-    """Test that the exclusion window is correctly initialized."""
-    img_dim = (50, 50)
-    win_rad = 10
-    estimator = NoiseEstimator(img_dim, win_rad)
-
-    mid_x, mid_y = img_dim[0] / 2, img_dim[1] / 2
-
-    for x in range(img_dim[0]):
-        for y in range(img_dim[1]):
-            # Pixel inside the exclusion radius should be False, others True
-            inside_radius = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2) <= win_rad
-            assert estimator.window[x, y] == (not inside_radius)
-
-
-def test_sigma_mad():
-    """Test the MAD-based standard deviation estimation."""
-    data = np.array(
-        [1, 1, 2, 2, 3, 3, 4, 4, 100]
-    )  # Outlier should not heavily influence MAD
-    expected_sigma = 1.4826 * np.median(np.abs(data - np.median(data)))
-
-    assert np.isclose(NoiseEstimator.sigma_mad(data), expected_sigma, atol=1e-4)
-
-
-def test_estimate_noise_without_default_window():
-    """Test noise estimation with the default exclusion window (no custom mask)."""
-    img_dim = (50, 50)
-    win_rad = 5
-    estimator = NoiseEstimator(img_dim, win_rad)
-
-    # Create a synthetic noisy image (Gaussian noise with mean=0, std=10)
-    np.random.seed(42)
-    image = np.random.normal(0, 10, img_dim)
-
-    noise_estimation = estimator.estimate_noise(image)
-
-    # The estimated noise should be close to 10 (the true std)
-    assert np.isclose(noise_estimation, 10, atol=2)
-
-
-def test_estimate_noise_with_custom_mask():
-    """Test noise estimation with a custom mask applied outside the exclusion radius."""
-    img_dim = (50, 50)
-    estimator = NoiseEstimator(img_dim, win_rad=5)
-
-    # Create synthetic noise with std=5
-    np.random.seed(42)
-    image = np.random.normal(0, 5, img_dim)
-
-    # Custom mask excluding top-left corner
-    custom_mask = np.ones(img_dim, dtype=bool)
-    custom_mask[:10, :10] = False  # Mask out top-left 10x10 pixels
-
-    noise_estimation = estimator.estimate_noise(image, mask=custom_mask)
-
-    assert np.isclose(noise_estimation, 5, atol=1)
-
-
-def test_apply_mask_with_none_mask():
-    """Test apply_mask when mask is None."""
-    img_dim = (10, 10)
-    estimator = NoiseEstimator(img_dim, win_rad=3)
-
-    result = estimator.apply_mask(None)  # Pass None as the mask
-
-    # It should return the window itself when no mask is provided
-    assert np.array_equal(
-        result, estimator.window
-    ), "apply_mask should return the window when mask is None."
-
-
-def test_apply_mask_with_valid_mask():
-    """Test apply_mask when a valid mask is provided."""
-    img_dim = (10, 10)
-    estimator = NoiseEstimator(img_dim, win_rad=3)
-
-    # Create a custom mask
-    custom_mask = np.ones(img_dim, dtype=bool)
-    custom_mask[5, 5] = False  # Set a pixel to False to exclude it from the window
-
-    result = estimator.apply_mask(custom_mask)
-
-    # Check that the mask was applied correctly: pixel (5, 5) should be False, others True
-    expected_result = estimator.window & custom_mask
-    assert np.array_equal(
-        result, expected_result
-    ), "apply_mask did not apply the mask correctly."
-
-
-def test_apply_mask_with_zeroed_mask():
-    """Test apply_mask when a zeroed mask is provided."""
-    img_dim = (10, 10)
-    estimator = NoiseEstimator(img_dim, win_rad=3)
-
-    # Create a mask where all pixels are excluded (False)
-    zeroed_mask = np.zeros(img_dim, dtype=bool)
-
-    result = estimator.apply_mask(zeroed_mask)
-
-    # The result should be an array of False values, as the mask excludes all pixels
-    expected_result = np.zeros(img_dim, dtype=bool)
-    assert np.array_equal(
-        result, expected_result
-    ), "apply_mask did not handle the zeroed mask correctly."
-
-
 def test_unobscured_zernike_projection():
     n_zernikes = 20
     wfe_dim = 256
diff --git a/src/wf_psf/training/train_utils.py b/src/wf_psf/training/train_utils.py
index 16deca0d..7365ade7 100644
--- a/src/wf_psf/training/train_utils.py
+++ b/src/wf_psf/training/train_utils.py
@@ -11,9 +11,10 @@
 
 import numpy as np
 import tensorflow as tf
-from typing import Optional, Callable, Union
+from typing import Optional, Callable
 from wf_psf.psf_models.psf_models import compile_PSF_model
-from wf_psf.utils.utils import NoiseEstimator, generalised_sigmoid
+from wf_psf.utils.noise import NoiseEstimator
+from wf_psf.utils.utils import generalised_sigmoid
 import logging
 
 logger = logging.getLogger(__name__)
@@ -367,31 +368,67 @@ def configure_optimizer_and_loss(
     return optimizer, loss, metrics
 
 
+def _resolve_training_outputs(
+    outputs: tf.Tensor, loss_name: Optional[str]
+) -> tuple[tf.Tensor, Optional[np.ndarray]]:
+    """
+    Resolve the training-specific ``outputs`` representation into a plain
+    image, and in case of a masked loss, a per-image mask.
+
+    Parameters
+    ----------
+    outputs: tf.Tensor
+        Image data. When ``loss_name`` starts with ``"masked_"``, a 4D tensor
+        of shape ``(batch_size, height, width, 2)`` is expected, where the last
+        dimension holds ``[image, mask]``. Otherwise, a 3D tensor of shape
+        ``(batch_size, height, width)`` is expected.
+    loss_name: str, optional
+        The name of the loss function used for training (e.g. ``loss.name``).
+        Any name starting with ``"masked_"`` (e.g. ``"masked_mean_squared_error"``)
+        is treated as a masked loss, so new masked losses work without updating
+        this function.
+
+    Returns
+    -------
+    tuple of (tf.Tensor, np.ndarray or None)
+        The resolved images and, for a masked loss, the corresponding boolean
+        masks (``outputs[..., 1]`` inverted); otherwise ``None``.
+    """
+    if loss_name is not None and loss_name.startswith("masked_"):
+        images = outputs[..., 0]
+        masks = np.array(1 - outputs[..., 1], dtype=bool)
+    else:
+        images = outputs
+        masks = None
+
+    return images, masks
+
+
 def calculate_sample_weights(
-    outputs: tf.Tensor,
-    use_sample_weights: bool,
-    loss: Union[str, Callable, None],
+    images: tf.Tensor,
+    masks: Optional[np.ndarray] = None,
     apply_sigmoid: bool = False,
     sigmoid_max_val: float = 5.0,
     sigmoid_power_k: float = 1.0,
-) -> Optional[np.ndarray]:
+) -> np.ndarray:
     """
     Calculate sample weights based on image noise standard deviation.
 
-    The function computes sample weights by estimating the noise standard deviation for each image, calculating the inverse variance,
-    and then normalizing the weights by dividing by the median.
+    The function estimates the per-observation noise standard deviation with
+    :class:`wf_psf.utils.noise.NoiseEstimator`, computes sample weights from the
+    inverse variance, and normalizes the weights by dividing by the median.
 
     Parameters
     ----------
-    outputs: tf.Tensor
+    images: tf.Tensor
         A 3D tensor of shape (batch_size, height, width) representing images, where the first dimension is the batch size
         and the next two dimensions are the image height and width.
-    use_sample_weights: bool
-        Flag indicating whether to compute sample weights. If True, sample weights will be computed based on the image noise.
-    loss: str, callable, optional
-        The loss function used for training. If the loss name is "masked_mean_squared_error", the function will calculate the noise standard deviation for masked images.
+    masks: np.ndarray, optional
+        A batch of boolean masks with the same shape as `images`, specifying
+        which pixels to include in the noise estimation for each image. If
+        None, only the exclusion window is used for every image. Default is None.
     apply_sigmoid: bool, optional
-        Flag indicating whether to apply a generalized sigmoid function to the sample weights. Default is True.
+        Flag indicating whether to apply a generalized sigmoid function to the sample weights. Default is False.
     sigmoid_max_val: float, optional
         The maximum value for the sigmoid function. Default is 5.0.
     sigmoid_power_k: float, optional
@@ -401,45 +438,27 @@ def calculate_sample_weights(
 
     Returns
     -------
-    np.ndarray or None
-        An array of sample weights, or None if `use_sample_weights` is False.
+    np.ndarray
+        An array of sample weights.
     """
-    if use_sample_weights:
-        img_dim = (outputs.shape[1], outputs.shape[2])
-        win_rad = np.ceil(outputs.shape[1] / 3.33)
-        std_est = NoiseEstimator(img_dim=img_dim, win_rad=win_rad)
-
-        if loss is not None and (
-            (isinstance(loss, str) and loss == "masked_mean_squared_error")
-            or (hasattr(loss, "name") and loss.name == "masked_mean_squared_error")
-        ):
-            
-            logger.info("Estimating noise standard deviation for masked images..")
-            images = outputs[..., 0]
-            masks = np.array(1 - outputs[..., 1], dtype=bool)
-            imgs_std = np.array(
-                [std_est.estimate_noise(_im, _win) for _im, _win in zip(images, masks)]
-            )
-        else:
-            logger.info("Estimating noise standard deviation for images..")
-            # Estimate noise standard deviation
-            imgs_std = np.array([std_est.estimate_noise(_im) for _im in outputs])
-
-        # Calculate variances
-        variances = imgs_std**2
-
-        # Use inverse variance for weights and scale by median
-        sample_weight = 1 / variances
-        sample_weight /= np.median(sample_weight)
-
-        # Apply a generalized sigmoid function to the sample weights
-        if apply_sigmoid:
-            sample_weight = generalised_sigmoid(
-                sample_weight, max_val=sigmoid_max_val, power_k=sigmoid_power_k
-            )
+    img_dim = (images.shape[1], images.shape[2])
+    std_est = NoiseEstimator(img_dim=img_dim)
 
-    else:
-        sample_weight = None
+    # Estimate the per-observation noise standard deviation
+    imgs_std = std_est.estimate_noise_batch(images, masks)
+
+    # Calculate variances
+    variances = imgs_std**2
+
+    # Use inverse variance for weights and scale by median
+    sample_weight = 1 / variances
+    sample_weight /= np.median(sample_weight)
+
+    # Apply a generalized sigmoid function to the sample weights
+    if apply_sigmoid:
+        sample_weight = generalised_sigmoid(
+            sample_weight, max_val=sigmoid_max_val, power_k=sigmoid_power_k
+        )
 
     return sample_weight
 
@@ -723,14 +742,19 @@ def general_train_cycle(
     )
 
     # Calculate sample weights
-    sample_weight = calculate_sample_weights(
-        outputs,
-        use_sample_weights,
-        loss,
-        apply_sigmoid,
-        sigmoid_max_val,
-        sigmoid_power_k,
-    )
+    if use_sample_weights:
+        images, masks = _resolve_training_outputs(
+            outputs, getattr(loss, "name", loss)
+        )
+        sample_weight = calculate_sample_weights(
+            images,
+            masks,
+            apply_sigmoid,
+            sigmoid_max_val,
+            sigmoid_power_k,
+        )
+    else:
+        sample_weight = None
 
     # Define the training cycle
     if cycle_def in ("parametric", "complete", "only-parametric"):
diff --git a/src/wf_psf/utils/noise.py b/src/wf_psf/utils/noise.py
new file mode 100644
index 00000000..b1d8b36c
--- /dev/null
+++ b/src/wf_psf/utils/noise.py
@@ -0,0 +1,203 @@
+"""Noise estimation utilities for image data.
+
+:Authors: Tobias Liaudat <tobias.liaudat@cea.fr>
+
+"""
+
+import numpy as np
+
+
+class NoiseEstimator:
+    """
+    A class for estimating noise levels in an image.
+
+    Parameters
+    ----------
+    img_dim : tuple of int
+        The dimensions of the image as (height, width).
+    win_rad : int, optional
+        The radius of the exclusion window (in pixels). If None, a default
+        radius is derived from `img_dim` (see :meth:`default_win_rad`).
+    """
+
+    def __init__(self, img_dim: tuple[int, int], win_rad: int = None) -> None:
+        """
+        Initialize a NoiseEstimator instance.
+
+        This constructor sets up the noise estimator by storing the image dimensions
+        and exclusion window radius, then initializes the exclusion window mask.
+
+        Parameters
+        ----------
+        img_dim : tuple of int
+            The dimensions of the image as (height, width).
+        win_rad : int, optional
+            The radius of the exclusion window in pixels. Pixels within this radius
+            of the image center are excluded from noise estimation. If None, the
+            radius is derived from `img_dim` via :meth:`default_win_rad`.
+
+        Notes
+        -----
+        The exclusion window is initialized automatically via _init_window(), creating
+        a boolean mask where pixels inside the exclusion radius are marked False
+        (excluded) and pixels outside are marked True (included).
+        """
+        self.img_dim: tuple[int, int] = img_dim
+        self.win_rad: int = (
+            win_rad if win_rad is not None else self.default_win_rad(img_dim)
+        )
+
+        self._init_window()  # Initialize self.window
+
+    @staticmethod
+    def default_win_rad(img_dim: tuple[int, int]) -> int:
+        """
+        Compute a default exclusion-window radius from the image dimensions.
+
+        Parameters
+        ----------
+        img_dim : tuple of int
+            The dimensions of the image as (height, width).
+
+        Returns
+        -------
+        int
+            The default exclusion-window radius, in pixels, for an image of
+            this size.
+        """
+        return int(np.ceil(img_dim[0] / 3.33))
+
+    def _init_window(self):
+        """
+        Initialize the exclusion window mask stored in self.window.
+
+        The mask is a boolean array of shape `self.img_dim` (rows, cols). Pixels
+        whose Euclidean distance from the image center is less than or equal to
+        `self.win_rad` are marked False (excluded); all other pixels are True
+        (included). The mask dtype is `bool`.
+
+        Notes
+        -----
+        - The image center is computed as (rows / 2, cols / 2). This yields a
+          floating-point center so the distance is computed with sub-pixel
+          precision; for even dimensions the center lies between pixels.
+        - The comparison uses "<=" so pixels exactly at distance `win_rad` are
+          excluded. Change to "<" if you prefer a strict interior exclusion.
+        - Time complexity is O(rows * cols) for mask construction.
+        - No return value; the constructed mask is assigned to `self.window`.
+        """
+        self.window = np.ones(self.img_dim, dtype=bool)
+
+        mid_x = self.img_dim[0] / 2
+        mid_y = self.img_dim[1] / 2
+
+        for _x in range(self.img_dim[0]):
+            for _y in range(self.img_dim[1]):
+                # If pixel is within the exclusion radius, set it to False
+                if np.sqrt((_x - mid_x) ** 2 + (_y - mid_y) ** 2) <= self.win_rad:
+                    self.window[_x, _y] = False
+
+    def apply_mask(self, mask: np.ndarray = None) -> np.ndarray:
+        """
+        Apply a given mask to the exclusion window.
+
+        Parameters
+        ----------
+        mask : np.ndarray, optional
+            A boolean mask to apply to the exclusion window. If None, the exclusion
+            window is returned without any modification.
+
+        Returns
+        -------
+        np.ndarray
+            The resulting boolean array after applying the mask to the exclusion window.
+        """
+        if mask is None:
+            return self.window  # Return just the window if no mask is provided
+        return self.window & mask  # Otherwise, apply the mask as usual
+
+    @staticmethod
+    def sigma_mad(x):
+        """
+        Robustly estimate the standard deviation using the Median Absolute Deviation (MAD).
+
+        Computes MAD = ``median(|x - median(x)|)`` and scales it by 1.4826 to make the
+        estimator consistent with the standard deviation for a Gaussian distribution:
+
+            sigma ≈ 1.4826 * MAD
+
+        Parameters
+        ----------
+        x : array-like
+            Input data. The values are flattened before computation. NaNs are not
+            specially handled and will propagate; remove or mask them prior to
+            calling if needed.
+
+        Returns
+        -------
+        float
+            Robust estimate of the standard deviation of the input data.
+
+        Notes
+        -----
+        - The MAD-based estimator is much less sensitive to outliers than the
+          sample standard deviation, making it appropriate for noisy data with
+          occasional large deviations.
+        - The constant 1.4826 is the scaling factor for consistency with the
+          standard deviation of a normal distribution.
+        """
+        return 1.4826 * np.median(np.abs(x - np.median(x)))
+
+    def estimate_noise(self, image: np.ndarray, mask: np.ndarray = None) -> float:
+        """
+        Estimates the noise level of an image using the MAD estimator.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            The input image for noise estimation.
+        mask : np.ndarray, optional
+            A boolean mask specifying which pixels to include in the noise estimation.
+            If None, the default exclusion window is used. The mask should have the same shape as `image`.
+
+        Returns
+        -------
+        float
+            The estimated noise standard deviation (MAD of the image pixels within the window or mask).
+        """
+        if mask is not None:
+            return self.sigma_mad(image[self.apply_mask(mask)])
+
+        # Use the default window if no mask is provided
+        return self.sigma_mad(image[self.window])
+
+    def estimate_noise_batch(
+        self, images: np.ndarray, masks: np.ndarray = None
+    ) -> np.ndarray:
+        """
+        Estimate the noise standard deviation for a batch of images.
+
+        Parameters
+        ----------
+        images : np.ndarray
+            A batch of images with shape ``(batch_size, height, width)``.
+        masks : np.ndarray, optional
+            A batch of boolean masks with the same shape as `images`, specifying
+            which pixels to include in the noise estimation for each image. If
+            None, only the exclusion window is used for every image.
+
+        Returns
+        -------
+        np.ndarray
+            A 1D array of shape ``(batch_size,)`` containing the estimated noise
+            standard deviation for each image.
+        """
+        if masks is None:
+            return np.array([self.estimate_noise(_im) for _im in images])
+
+        return np.array(
+            [
+                self.estimate_noise(_im, _mask)
+                for _im, _mask in zip(images, masks)
+            ]
+        )
diff --git a/src/wf_psf/utils/utils.py b/src/wf_psf/utils/utils.py
index 17219ad9..401cc6d8 100644
--- a/src/wf_psf/utils/utils.py
+++ b/src/wf_psf/utils/utils.py
@@ -437,149 +437,6 @@ def add_noise(image, desired_SNR):
     return noisy_image
 
 
-class NoiseEstimator:
-    """
-    A class for estimating noise levels in an image.
-
-    Parameters
-    ----------
-    img_dim : tuple of int
-        The dimensions of the image as (height, width).
-    win_rad : int
-        The radius of the exclusion window (in pixels).
-    """
-
-    def __init__(self, img_dim: tuple[int, int], win_rad: int) -> None:
-        """
-        Initialize a NoiseEstimator instance.
-
-        This constructor sets up the noise estimator by storing the image dimensions
-        and exclusion window radius, then initializes the exclusion window mask.
-
-        Parameters
-        ----------
-        img_dim : tuple of int
-            The dimensions of the image as (height, width).
-        win_rad : int
-            The radius of the exclusion window in pixels. Pixels within this radius
-            of the image center are excluded from noise estimation.
-
-        Notes
-        -----
-        The exclusion window is initialized automatically via _init_window(), creating
-        a boolean mask where pixels inside the exclusion radius are marked False
-        (excluded) and pixels outside are marked True (included).
-        """
-        self.img_dim: tuple[int, int] = img_dim
-        self.win_rad: int = win_rad
-
-        self._init_window()  # Initialize self.window
-
-    def _init_window(self):
-        """
-        Initialize the exclusion window mask stored in self.window.
-
-        The mask is a boolean array of shape `self.img_dim` (rows, cols). Pixels
-        whose Euclidean distance from the image center is less than or equal to
-        `self.win_rad` are marked False (excluded); all other pixels are True
-        (included). The mask dtype is `bool`.
-
-        Notes
-        -----
-        - The image center is computed as (rows / 2, cols / 2). This yields a
-          floating-point center so the distance is computed with sub-pixel
-          precision; for even dimensions the center lies between pixels.
-        - The comparison uses "<=" so pixels exactly at distance `win_rad` are
-          excluded. Change to "<" if you prefer a strict interior exclusion.
-        - Time complexity is O(rows * cols) for mask construction.
-        - No return value; the constructed mask is assigned to `self.window`.
-        """
-        self.window = np.ones(self.img_dim, dtype=bool)
-
-        mid_x = self.img_dim[0] / 2
-        mid_y = self.img_dim[1] / 2
-
-        for _x in range(self.img_dim[0]):
-            for _y in range(self.img_dim[1]):
-                # If pixel is within the exclusion radius, set it to False
-                if np.sqrt((_x - mid_x) ** 2 + (_y - mid_y) ** 2) <= self.win_rad:
-                    self.window[_x, _y] = False
-
-    def apply_mask(self, mask: np.ndarray = None) -> np.ndarray:
-        """
-        Apply a given mask to the exclusion window.
-
-        Parameters
-        ----------
-        mask : np.ndarray, optional
-            A boolean mask to apply to the exclusion window. If None, the exclusion
-            window is returned without any modification.
-
-        Returns
-        -------
-        np.ndarray
-            The resulting boolean array after applying the mask to the exclusion window.
-        """
-        if mask is None:
-            return self.window  # Return just the window if no mask is provided
-        return self.window & mask  # Otherwise, apply the mask as usual
-
-    @staticmethod
-    def sigma_mad(x):
-        """
-        Robustly estimate the standard deviation using the Median Absolute Deviation (MAD).
-
-        Computes MAD = ``median(|x - median(x)|)`` and scales it by 1.4826 to make the
-        estimator consistent with the standard deviation for a Gaussian distribution:
-
-            sigma ≈ 1.4826 * MAD
-
-        Parameters
-        ----------
-        x : array-like
-            Input data. The values are flattened before computation. NaNs are not
-            specially handled and will propagate; remove or mask them prior to
-            calling if needed.
-
-        Returns
-        -------
-        float
-            Robust estimate of the standard deviation of the input data.
-
-        Notes
-        -----
-        - The MAD-based estimator is much less sensitive to outliers than the
-          sample standard deviation, making it appropriate for noisy data with
-          occasional large deviations.
-        - The constant 1.4826 is the scaling factor for consistency with the
-          standard deviation of a normal distribution.
-        """
-        return 1.4826 * np.median(np.abs(x - np.median(x)))
-
-    def estimate_noise(self, image: np.ndarray, mask: np.ndarray = None) -> float:
-        """
-        Estimates the noise level of an image using the MAD estimator.
-
-        Parameters
-        ----------
-        image : np.ndarray
-            The input image for noise estimation.
-        mask : np.ndarray, optional
-            A boolean mask specifying which pixels to include in the noise estimation.
-            If None, the default exclusion window is used. The mask should have the same shape as `image`.
-
-        Returns
-        -------
-        float
-            The estimated noise standard deviation (MAD of the image pixels within the window or mask).
-        """
-        if mask is not None:
-            return self.sigma_mad(image[self.apply_mask(mask)])
-
-        # Use the default window if no mask is provided
-        return self.sigma_mad(image[self.window])
-
-
 class ZernikeInterpolation:
     """Interpolate Zernike coefficients using K-nearest RBF splines.