Improve confusion matrix plotting (#2358)

Round floats to avoid floating point errors leading to UI overflow.
Remove overlapping text in multilabel plots
by reducing redundant `Predicted class` and `True class` labels.
Use `constrained_layout` to prevent some text from being cut off.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Nicki Skafte Detlefsen <skaftenicki@gmail.com>

CHANGELOG.md

@@ -47,6 +47,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fixed initialize aggregation metrics with default floating type ([#2366](https://github.com/Lightning-AI/torchmetrics/pull/2366))
 
+
+- Fixed plotting of confusion matrices ([#2358](https://github.com/Lightning-AI/torchmetrics/pull/2358))
+
 ---
 
 ## [1.3.0] - 2024-01-10

src/torchmetrics/utilities/plot.py

@@ -242,15 +242,17 @@ def plot_confusion_matrix(
  fig_label = None
  labels = labels or np.arange(n_classes).tolist()
 
- fig, axs = plt.subplots(nrows=rows, ncols=cols) if ax is None else (ax.get_figure(), ax)
+ fig, axs = plt.subplots(nrows=rows, ncols=cols, constrained_layout=True) if ax is None else (ax.get_figure(), ax)
  axs = trim_axs(axs, nb)
  for i in range(nb):
  ax = axs[i] if rows != 1 and cols != 1 else axs
  if fig_label is not None:
  ax.set_title(f"Label {fig_label[i]}", fontsize=15)
  ax.imshow(confmat[i].cpu().detach() if confmat.ndim == 3 else confmat.cpu().detach())
- ax.set_xlabel("Predicted class", fontsize=15)
- ax.set_ylabel("True class", fontsize=15)
+ if i // cols == rows - 1: # bottom row only
+ ax.set_xlabel("Predicted class", fontsize=15)
+ if i % cols == 0: # leftmost column only
+ ax.set_ylabel("True class", fontsize=15)
  ax.set_xticks(list(range(n_classes)))
  ax.set_yticks(list(range(n_classes)))
  ax.set_xticklabels(labels, rotation=45, fontsize=10)
@@ -259,7 +261,7 @@ def plot_confusion_matrix(
  if add_text:
  for ii, jj in product(range(n_classes), range(n_classes)):
  val = confmat[i, ii, jj] if confmat.ndim == 3 else confmat[ii, jj]
- ax.text(jj, ii, str(val.item()), ha="center", va="center", fontsize=15)
+ ax.text(jj, ii, str(round(val.item(), 2)), ha="center", va="center", fontsize=15)
 
  return fig, axs