Add new metrics: SAM

docs/source/references/functional.rst

@@ -281,7 +281,7 @@ universal_image_quality_index [func]
 universal_spectral_angle_mapper [func]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: torchmetrics.functional.universal_spectral_angle_mapper
+.. autofunction:: torchmetrics.functional.spectral_angle_mapper
  :noindex:
 
 

tests/image/test_sam.py

@@ -19,7 +19,7 @@
 
 from tests.helpers import seed_all
 from tests.helpers.testers import BATCH_SIZE, NUM_BATCHES, MetricTester
-from torchmetrics.functional.image.sam import universal_spectral_angle_mapper
+from torchmetrics.functional.image.sam import spectral_angle_mapper
 from torchmetrics.image.sam import SpectralAngleMapper
 
 seed_all(42)
@@ -28,9 +28,9 @@
 
 _inputs = []
 for size, channel, dtype in [
- (12, 1, torch.float),
+ (12, 3, torch.float),
  (13, 3, torch.float32),
- (14, 1, torch.double),
+ (14, 3, torch.double),
  (15, 3, torch.float64),
 ]:
  preds = torch.rand(NUM_BATCHES, BATCH_SIZE, channel, size, size, dtype=dtype)
@@ -42,7 +42,9 @@ def _sk_sam(preds, target, reduction):
  # reshape to (batch_size, channel, height*width)
  B, C, H, W = preds.shape
  sk_preds = preds.reshape(B, C, H * W)
+ sk_preds = torch.movedim(sk_preds, 1, -1)
  sk_target = target.reshape(B, C, H * W)
+ sk_target = torch.movedim(sk_target, 1, -1)
  # compute arccos of cosine similarity
  dot_product = (sk_preds * sk_target).sum(dim=-1)
  preds_norm = sk_preds.norm(dim=-1)
@@ -82,7 +84,7 @@ def test_sam_functional(self, reduction, preds, target):
  self.run_functional_metric_test(
  preds,
  target,
- universal_spectral_angle_mapper,
+ spectral_angle_mapper,
  partial(_sk_sam, reduction=reduction),
  metric_args=dict(reduction=reduction),
  )
@@ -94,12 +96,12 @@ def test_sam_half_cpu(self, reduction, preds, target):
  preds,
  target,
  SpectralAngleMapper,
- universal_spectral_angle_mapper,
+ spectral_angle_mapper,
  )
 
  @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
  def test_sam_half_gpu(self, reduction, preds, target):
- self.run_precision_test_gpu(preds, target, SpectralAngleMapper, universal_spectral_angle_mapper)
+ self.run_precision_test_gpu(preds, target, SpectralAngleMapper, spectral_angle_mapper)
 
 
 def test_error_on_different_shape(metric_class=SpectralAngleMapper):
@@ -118,3 +120,9 @@ def test_error_on_invalid_type(metric_class=SpectralAngleMapper):
  metric = metric_class()
  with pytest.raises(TypeError):
  metric(torch.randn([3, 16, 16]), torch.randn([3, 16, 16], dtype=torch.float64))
+
+
+def test_error_on_grayscale_image(metric_class=SpectralAngleMapper):
+ metric = metric_class()
+ with pytest.raises(ValueError):
+ metric(torch.randn([16, 1, 16, 16]), torch.randn([16, 1, 16, 16]))

torchmetrics/functional/__init__.py

@@ -35,7 +35,7 @@
 from torchmetrics.functional.classification.stat_scores import stat_scores
 from torchmetrics.functional.image.gradients import image_gradients
 from torchmetrics.functional.image.psnr import peak_signal_noise_ratio
-from torchmetrics.functional.image.sam import universal_spectral_angle_mapper
+from torchmetrics.functional.image.sam import spectral_angle_mapper
 from torchmetrics.functional.image.ssim import (
  multiscale_structural_similarity_index_measure,
  structural_similarity_index_measure,
@@ -145,7 +145,7 @@
  "symmetric_mean_absolute_percentage_error",
  "translation_edit_rate",
  "universal_image_quality_index",
- "universal_spectral_angle_mapper",
+ "spectral_angle_mapper",
  "word_error_rate",
  "char_error_rate",
  "match_error_rate",

torchmetrics/functional/image/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from torchmetrics.functional.image.gradients import image_gradients # noqa: F401
 from torchmetrics.functional.image.psnr import peak_signal_noise_ratio # noqa: F401
+from torchmetrics.functional.image.sam import spectral_angle_mapper # noqa: F401
 from torchmetrics.functional.image.ssim import ( # noqa: F401
  multiscale_structural_similarity_index_measure,
  structural_similarity_index_measure,

torchmetrics/functional/image/sam.py

@@ -42,6 +42,11 @@ def _sam_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
  "Expected `preds` and `target` to have BxCxHxW shape."
  f" Got preds: {preds.shape} and target: {target.shape}."
  )
+ if (preds.shape[1] <= 1) or (target.shape[1] <= 1):
+ raise ValueError(
+ "Expected dimension `preds` and `target` is larger than 1."
+ f" Got preds: {preds.shape[1]} and target: {target.shape[1]}."
+ )
  return preds, target
 
 
@@ -62,20 +67,22 @@ def _sam_compute(
  - ``'none'``: no reduction will be applied
 
  Example:
- >>> preds = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(42))
- >>> target = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(123))
+ >>> preds = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(42))
+ >>> target = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(123))
  >>> preds, target = _sam_update(preds, target)
  >>> _sam_compute(preds, target)
- tensor(0.7377)
+ tensor(0.5943)
  """
  B, C, H, W = preds.shape
  preds = preds.reshape(B, C, H * W)
+ preds = torch.movedim(preds, 1, -1)
  target = target.reshape(B, C, H * W)
+ target = torch.movedim(target, 1, -1)
  sam_score = torch.clip(cosine_similarity(preds, target, reduction="none"), -1, 1).arccos()
  return reduce(sam_score, reduction)
 
 
-def universal_spectral_angle_mapper(
+def spectral_angle_mapper(
  preds: Tensor,
  target: Tensor,
  reduction: Literal["elementwise_mean", "sum", "none"] = "elementwise_mean",
@@ -101,11 +108,11 @@ def universal_spectral_angle_mapper(
  If ``preds`` and ``target`` don't have ``BxCxHxW shape``.
 
  Example:
- >>> from torchmetrics.functional import universal_spectral_angle_mapper
- >>> preds = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(42))
- >>> target = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(123))
- >>> universal_spectral_angle_mapper(preds, target)
- tensor(0.7377)
+ >>> from torchmetrics.functional import spectral_angle_mapper
+ >>> preds = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(42))
+ >>> target = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(123))
+ >>> spectral_angle_mapper(preds, target)
+ tensor(0.5943)
 
  References: Roberta H. Yuhas, Alexander F. H. Goetz and Joe W. Boardman, "Discrimination among semi-arid
  landscape endmembers using the Spectral Angle Mapper (SAM) algorithm" in PL, Summaries of the Third Annual JPL

torchmetrics/image/sam.py

@@ -43,11 +43,11 @@ class SpectralAngleMapper(Metric):
  Example:
  >>> import torch
  >>> from torchmetrics import SpectralAngleMapper
- >>> preds = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(42))
- >>> target = torch.rand([16, 1, 16, 16], generator=torch.manual_seed(123))
+ >>> preds = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(42))
+ >>> target = torch.rand([16, 3, 16, 16], generator=torch.manual_seed(123))
  >>> sam = SpectralAngleMapper()
  >>> sam(preds, target)
- tensor(0.7377)
+ tensor(0.5943)
  """
 
  preds: List[Tensor]