Add plotting 8/n

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)