Add plotting 8/n (#1600)

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  [#1581](https://github.com/Lightning-AI/metrics/pull/1581),
  [#1585](https://github.com/Lightning-AI/metrics/pull/1585),
  [#1593](https://github.com/Lightning-AI/metrics/pull/1593),
+ [#1600](https://github.com/Lightning-AI/metrics/pull/1600),
 )
 
 

requirements/docs.txt

@@ -16,3 +16,4 @@ sphinx-copybutton>=0.3
 -r visual.txt
 -r audio.txt
 -r detection.txt
+-r image.txt

src/torchmetrics/image/fid.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from copy import deepcopy
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional, Sequence, Union
 
 import numpy as np
 import torch
@@ -22,7 +22,11 @@
 
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_info
-from torchmetrics.utilities.imports import _SCIPY_AVAILABLE, _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _SCIPY_AVAILABLE, _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+
+if not _MATPLOTLIB_AVAILABLE:
+ __doctest_skip__ = ["FrechetInceptionDistance.plot"]
 
 if _TORCH_FIDELITY_AVAILABLE:
  from torch_fidelity.feature_extractor_inceptionv3 import FeatureExtractorInceptionV3 as _FeatureExtractorInceptionV3
@@ -31,7 +35,7 @@
  class _FeatureExtractorInceptionV3(Module):
  pass
 
- __doctest_skip__ = ["FrechetInceptionDistance", "FID"]
+ __doctest_skip__ = ["FrechetInceptionDistance", "FrechetInceptionDistance.plot"]
 
 
 if _SCIPY_AVAILABLE:
@@ -303,3 +307,54 @@ def reset(self) -> None:
  self.real_features_num_samples = real_features_num_samples
  else:
  super().reset()
+
+ def plot(
+ self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
+ ) -> _PLOT_OUT_TYPE:
+ """Plot a single or multiple values from the metric.
+
+ Args:
+ val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
+ If no value is provided, will automatically call `metric.compute` and plot that result.
+ ax: An matplotlib axis object. If provided will add plot to that axis
+
+ Returns:
+ Figure and Axes object
+
+ Raises:
+ ModuleNotFoundError:
+ If `matplotlib` is not installed
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting a single value
+ >>> import torch
+ >>> from torchmetrics.image.fid import FrechetInceptionDistance
+ >>> imgs_dist1 = torch.randint(0, 200, (100, 3, 299, 299), dtype=torch.uint8)
+ >>> imgs_dist2 = torch.randint(100, 255, (100, 3, 299, 299), dtype=torch.uint8)
+ >>> metric = FrechetInceptionDistance(feature=64)
+ >>> metric.update(imgs_dist1, real=True)
+ >>> metric.update(imgs_dist2, real=False)
+ >>> fig_, ax_ = metric.plot()
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting multiple values
+ >>> import torch
+ >>> from torchmetrics.image.fid import FrechetInceptionDistance
+ >>> imgs_dist1 = lambda: torch.randint(0, 200, (100, 3, 299, 299), dtype=torch.uint8)
+ >>> imgs_dist2 = lambda: torch.randint(100, 255, (100, 3, 299, 299), dtype=torch.uint8)
+ >>> metric = FrechetInceptionDistance(feature=64)
+ >>> values = [ ]
+ >>> for _ in range(3):
+ ... metric.update(imgs_dist1(), real=True)
+ ... metric.update(imgs_dist2(), real=False)
+ ... values.append(metric.compute())
+ ... metric.reset()
+ >>> fig_, ax_ = metric.plot(values)
+
+
+ """
+ return self._plot(val, ax)

src/torchmetrics/image/inception.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, List, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -21,9 +21,14 @@
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_warn
 from torchmetrics.utilities.data import dim_zero_cat
-from torchmetrics.utilities.imports import _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
 
-__doctest_requires__ = {("InceptionScore", "IS"): ["torch_fidelity"]}
+if not _MATPLOTLIB_AVAILABLE:
+ __doctest_skip__ = ["InceptionScore.plot"]
+
+
+__doctest_requires__ = {("InceptionScore", "InceptionScore.plot"): ["torch_fidelity"]}
 
 
 class InceptionScore(Metric):
@@ -162,3 +167,46 @@ def compute(self) -> Tuple[Tensor, Tensor]:
 
  # return mean and std
  return kl.mean(), kl.std()
+
+ def plot(
+ self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
+ ) -> _PLOT_OUT_TYPE:
+ """Plot a single or multiple values from the metric.
+
+ Args:
+ val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
+ If no value is provided, will automatically call `metric.compute` and plot that result.
+ ax: An matplotlib axis object. If provided will add plot to that axis
+
+ Returns:
+ Figure and Axes object
+
+ Raises:
+ ModuleNotFoundError:
+ If `matplotlib` is not installed
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting a single value
+ >>> import torch
+ >>> from torchmetrics.image.inception import InceptionScore
+ >>> metric = InceptionScore()
+ >>> metric.update(torch.randint(0, 255, (50, 3, 299, 299), dtype=torch.uint8))
+ >>> fig_, ax_ = metric.plot() # the returned plot only shows the mean value by default
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting multiple values
+ >>> import torch
+ >>> from torchmetrics.image.inception import InceptionScore
+ >>> metric = InceptionScore()
+ >>> values = [ ]
+ >>> for _ in range(3):
+ ... # we index by 0 such that only the mean value is plotted
+ ... values.append(metric(torch.randint(0, 255, (50, 3, 299, 299), dtype=torch.uint8))[0])
+ >>> fig_, ax_ = metric.plot(values)
+ """
+ val = val or self.compute()[0] # by default we select the mean to plot
+ return self._plot(val, ax)

src/torchmetrics/image/kid.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -21,9 +21,13 @@
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_warn
 from torchmetrics.utilities.data import dim_zero_cat
-from torchmetrics.utilities.imports import _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _TORCH_FIDELITY_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
 
-__doctest_requires__ = {("KernelInceptionDistance", "KID"): ["torch_fidelity"]}
+if not _MATPLOTLIB_AVAILABLE:
+ __doctest_skip__ = ["KernelInceptionDistance.plot"]
+
+__doctest_requires__ = {("KernelInceptionDistance", "KernelInceptionDistance.plot"): ["torch_fidelity"]}
 
 
 def maximum_mean_discrepancy(k_xx: Tensor, k_xy: Tensor, k_yy: Tensor) -> Tensor:
@@ -274,3 +278,53 @@ def reset(self) -> None:
  self._defaults["real_features"] = value
  else:
  super().reset()
+
+ def plot(
+ self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
+ ) -> _PLOT_OUT_TYPE:
+ """Plot a single or multiple values from the metric.
+
+ Args:
+ val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
+ If no value is provided, will automatically call `metric.compute` and plot that result.
+ ax: An matplotlib axis object. If provided will add plot to that axis
+
+ Returns:
+ Figure and Axes object
+
+ Raises:
+ ModuleNotFoundError:
+ If `matplotlib` is not installed
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting a single value
+ >>> import torch
+ >>> from torchmetrics.image.kid import KernelInceptionDistance
+ >>> imgs_dist1 = torch.randint(0, 200, (30, 3, 299, 299), dtype=torch.uint8)
+ >>> imgs_dist2 = torch.randint(100, 255, (30, 3, 299, 299), dtype=torch.uint8)
+ >>> metric = KernelInceptionDistance(subsets=3, subset_size=20)
+ >>> metric.update(imgs_dist1, real=True)
+ >>> metric.update(imgs_dist2, real=False)
+ >>> fig_, ax_ = metric.plot()
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting multiple values
+ >>> import torch
+ >>> from torchmetrics.image.kid import KernelInceptionDistance
+ >>> imgs_dist1 = lambda: torch.randint(0, 200, (30, 3, 299, 299), dtype=torch.uint8)
+ >>> imgs_dist2 = lambda: torch.randint(100, 255, (30, 3, 299, 299), dtype=torch.uint8)
+ >>> metric = KernelInceptionDistance(subsets=3, subset_size=20)
+ >>> values = [ ]
+ >>> for _ in range(3):
+ ... metric.update(imgs_dist1(), real=True)
+ ... metric.update(imgs_dist2(), real=False)
+ ... values.append(metric.compute()[0])
+ ... metric.reset()
+ >>> fig_, ax_ = metric.plot(values)
+ """
+ val = val or self.compute()[0] # by default we select the mean to plot
+ return self._plot(val, ax)

src/torchmetrics/image/lpip.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-from typing import Any, List
+from typing import Any, List, Optional, Sequence, Union
 
 import torch
 from torch import Tensor
@@ -21,7 +21,11 @@
 
 from torchmetrics.metric import Metric
 from torchmetrics.utilities.checks import _SKIP_SLOW_DOCTEST, _try_proceed_with_timeout
-from torchmetrics.utilities.imports import _LPIPS_AVAILABLE
+from torchmetrics.utilities.imports import _LPIPS_AVAILABLE, _MATPLOTLIB_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+
+if not _MATPLOTLIB_AVAILABLE:
+ __doctest_skip__ = ["LearnedPerceptualImagePatchSimilarity.plot"]
 
 if _LPIPS_AVAILABLE:
  from lpips import LPIPS as _LPIPS
@@ -30,13 +34,13 @@ def _download_lpips() -> None:
  _LPIPS(pretrained=True, net="vgg")
 
  if _SKIP_SLOW_DOCTEST and not _try_proceed_with_timeout(_download_lpips):
- __doctest_skip__ = ["LearnedPerceptualImagePatchSimilarity", "LPIPS"]
+ __doctest_skip__ = ["LearnedPerceptualImagePatchSimilarity", "LearnedPerceptualImagePatchSimilarity.plot"]
 else:
 
  class _LPIPS(Module):
  pass
 
- __doctest_skip__ = ["LearnedPerceptualImagePatchSimilarity", "LPIPS"]
+ __doctest_skip__ = ["LearnedPerceptualImagePatchSimilarity", "LearnedPerceptualImagePatchSimilarity.plot"]
 
 
 class NoTrainLpips(_LPIPS):
@@ -167,3 +171,44 @@ def compute(self) -> Tensor:
  if self.reduction == "sum":
  return self.sum_scores
  return None
+
+ def plot(
+ self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
+ ) -> _PLOT_OUT_TYPE:
+ """Plot a single or multiple values from the metric.
+
+ Args:
+ val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
+ If no value is provided, will automatically call `metric.compute` and plot that result.
+ ax: An matplotlib axis object. If provided will add plot to that axis
+
+ Returns:
+ Figure and Axes object
+
+ Raises:
+ ModuleNotFoundError:
+ If `matplotlib` is not installed
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting a single value
+ >>> import torch
+ >>> from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+ >>> metric = LearnedPerceptualImagePatchSimilarity()
+ >>> metric.update(torch.rand(10, 3, 100, 100), torch.rand(10, 3, 100, 100))
+ >>> fig_, ax_ = metric.plot()
+
+ .. plot::
+ :scale: 75
+
+ >>> # Example plotting multiple values
+ >>> import torch
+ >>> from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+ >>> metric = LearnedPerceptualImagePatchSimilarity()
+ >>> values = [ ]
+ >>> for _ in range(3):
+ ... values.append(metric(torch.rand(10, 3, 100, 100), torch.rand(10, 3, 100, 100)))
+ >>> fig_, ax_ = metric.plot(values)
+ """
+ return self._plot(val, ax)