Clamp variance calculation in certain image metrics

CHANGELOG.md

@@ -53,6 +53,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fixed plotting of confusion matrices ([#2358](https://github.com/Lightning-AI/torchmetrics/pull/2358))
 
+
+- Fixed negative variance estimates in certain image metrics ([#2378](https://github.com/Lightning-AI/torchmetrics/pull/2378))
+
 ---
 
 ## [1.3.0] - 2024-01-10

src/torchmetrics/functional/image/ssim.py

@@ -154,9 +154,10 @@ def _ssim_update(
     mu_target_sq = output_list[1].pow(2)
     mu_pred_target = output_list[0] * output_list[1]
 
-    sigma_pred_sq = output_list[2] - mu_pred_sq
-    sigma_target_sq = output_list[3] - mu_target_sq
-    sigma_pred_target = output_list[4] - mu_pred_target
+    # Calculate the variance of the predicted and target images, should be non-negative
+    sigma_pred_sq = torch.clamp(output_list[2] - mu_pred_sq, min=0.0)
+    sigma_target_sq = torch.clamp(output_list[3] - mu_target_sq, min=0.0)
+    sigma_pred_target = torch.clamp(output_list[4] - mu_pred_target, min=0.0)
 
     upper = 2 * sigma_pred_target.to(dtype) + c2
     lower = (sigma_pred_sq + sigma_target_sq).to(dtype) + c2

src/torchmetrics/functional/image/uqi.py

@@ -102,9 +102,10 @@ def _uqi_compute(
     mu_target_sq = output_list[1].pow(2)
     mu_pred_target = output_list[0] * output_list[1]
 
-    sigma_pred_sq = output_list[2] - mu_pred_sq
-    sigma_target_sq = output_list[3] - mu_target_sq
-    sigma_pred_target = output_list[4] - mu_pred_target
+    # Calculate the variance of the predicted and target images, should be non-negative
+    sigma_pred_sq = torch.clamp(output_list[2] - mu_pred_sq, min=0.0)
+    sigma_target_sq = torch.clamp(output_list[3] - mu_target_sq, min=0.0)
+    sigma_pred_target = torch.clamp(output_list[4] - mu_pred_target, min=0.0)
 
     upper = 2 * sigma_pred_target
     lower = sigma_pred_sq + sigma_target_sq