Add _repr_html_ for subject too

Author: fepegar

src/torchio/data/image.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import base64
-import io
 import warnings
 from collections import Counter
 from collections.abc import Callable
@@ -203,25 +201,16 @@ def __repr__(self):
 
     def _repr_html_(self):
         try:
-            from matplotlib import pyplot as plt
             from matplotlib.figure import Figure
         except ImportError:
             return self.__repr__()
 
-        buffer = io.BytesIO()
-        fig = self.plot(
-            return_fig=True,
-            output_path=buffer,
-            show=False,
-            savefig_kwargs={'bbox_inches': 'tight'},
-        )
+        fig = self.plot(return_fig=True, show=False)
         assert isinstance(fig, Figure)
-        plt.close(fig)
-        buffer.seek(0)
 
-        img_str = base64.b64encode(buffer.read()).decode('utf-8')
-        html = f'<img src="data:image/png;base64,{img_str}"/>'
-        return html
+        from ..visualization import _figure_to_html
+
+        return _figure_to_html(fig)
 
     def __getitem__(self, item):
         if isinstance(item, (slice, int, tuple)):

src/torchio/data/subject.py

@@ -14,6 +14,8 @@
 from .image import Image
 
 if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
     from ..transforms import Compose
     from ..transforms import Transform
 
@@ -66,6 +68,19 @@ def __repr__(self):
         )
         return string
 
+    def _repr_html_(self):
+        try:
+            from matplotlib.figure import Figure
+        except ImportError:
+            return self.__repr__()
+
+        fig = self.plot(return_fig=True, show=False)
+        assert isinstance(fig, Figure)
+
+        from ..visualization import _figure_to_html
+
+        return _figure_to_html(fig)
+
     def __len__(self):
         return len(self.get_images(intensity_only=False))
 
@@ -412,13 +427,16 @@ def remove_image(self, image_name: str) -> None:
         del self[image_name]
         delattr(self, image_name)
 
-    def plot(self, **kwargs) -> None:
+    def plot(self, return_fig: bool = False, **kwargs) -> None | Figure:
         """Plot images using matplotlib.
 
         Args:
+            return_fig: If ``True``, return the figure instead of showing it.
             **kwargs: Keyword arguments that will be passed on to
                 [`plot()`][torchio.Image.plot].
         """
         from ..visualization import plot_subject  # avoid circular import
 
-        plot_subject(self, **kwargs)
+        figure = plot_subject(self, **kwargs)
+        if return_fig:
+            return figure

src/torchio/visualization.py

@@ -38,6 +38,21 @@ def import_mpl_plt():
     return mpl, plt
 
 
+def _figure_to_html(fig: Figure) -> str:
+    """Convert a matplotlib Figure to an HTML img tag with base64-encoded PNG."""
+    import base64
+    import io
+
+    from matplotlib import pyplot as plt
+
+    buffer = io.BytesIO()
+    fig.savefig(buffer, format='png', bbox_inches='tight')
+    plt.close(fig)
+    buffer.seek(0)
+    img_str = base64.b64encode(buffer.read()).decode('utf-8')
+    return f'<img src="data:image/png;base64,{img_str}"/>'
+
+
 def rotate(image: np.ndarray, *, radiological: bool = True, n: int = -1) -> np.ndarray:
     # Rotate for visualization purposes
     image = np.rot90(image, n, axes=(0, 1))
@@ -214,8 +229,9 @@ def plot_subject(
     output_path=None,
     figsize=None,
     clear_axes=True,
+    savefig_kwargs: dict[str, Any] | None = None,
     **plot_volume_kwargs,
-):
+) -> Figure:
     _, plt = import_mpl_plt()
     num_images = len(subject)
     many_images = num_images > 2
@@ -252,9 +268,12 @@ def plot_subject(
             axis.set_title(f'{name} ({axis_name})')
     plt.tight_layout()
     if output_path is not None:
-        fig.savefig(output_path)
+        if savefig_kwargs is None:
+            savefig_kwargs = {}
+        fig.savefig(output_path, **savefig_kwargs)
     if show:
         plt.show()
+    return fig
 
 
 def get_num_bins(x: np.ndarray) -> int: