Reduce matplotlib multithreading errors with Agg backend

jamesdolezal · jamesdolezal · commit 67db53bfde25 · 2024-08-13T18:59:56.000-04:00
diff --git a/slideflow/heatmap.py b/slideflow/heatmap.py
@@ -740,70 +740,71 @@ def save(
         """
         import matplotlib.pyplot as plt
 
-        if self.predictions is None:
-            raise errors.HeatmapError(
-                "Cannot plot Heatmap which is not yet generated; generate with "
-                "either heatmap.generate() or Heatmap(..., generate=True)"
-            )
+        with sf.util.matplotlib_backend('Agg'):
+            if self.predictions is None:
+                raise errors.HeatmapError(
+                    "Cannot plot Heatmap which is not yet generated; generate with "
+                    "either heatmap.generate() or Heatmap(..., generate=True)"
+                )
 
-        # Save heatmaps in .npz format
-        self.save_npz(os.path.join(outdir, f'{self.slide.name}.npz'))
+            # Save heatmaps in .npz format
+            self.save_npz(os.path.join(outdir, f'{self.slide.name}.npz'))
 
-        def _savefig(label, bbox_inches='tight', **kwargs):
-            plt.savefig(
-                os.path.join(outdir, f'{self.slide.name}-{label}.png'),
-                bbox_inches=bbox_inches,
-                **kwargs
-            )
+            def _savefig(label, bbox_inches='tight', **kwargs):
+                plt.savefig(
+                    os.path.join(outdir, f'{self.slide.name}-{label}.png'),
+                    bbox_inches=bbox_inches,
+                    **kwargs
+                )
 
-        log.info('Saving base figures...')
+            log.info('Saving base figures...')
 
-        # Prepare matplotlib figure
-        ax = self._prepare_ax()
+            # Prepare matplotlib figure
+            ax = self._prepare_ax()
 
-        thumb_kwargs = dict(roi_color=roi_color, linewidth=linewidth)
+            thumb_kwargs = dict(roi_color=roi_color, linewidth=linewidth)
 
-        # Save base thumbnail as separate figure
-        self.plot_thumbnail(show_roi=False, ax=ax, **thumb_kwargs)  # type: ignore
-        _savefig('raw')
+            # Save base thumbnail as separate figure
+            self.plot_thumbnail(show_roi=False, ax=ax, **thumb_kwargs)  # type: ignore
+            _savefig('raw')
 
-        # Save thumbnail + ROI as separate figure
-        self.plot_thumbnail(show_roi=True, ax=ax, **thumb_kwargs)  # type: ignore
-        _savefig('raw+roi')
+            # Save thumbnail + ROI as separate figure
+            self.plot_thumbnail(show_roi=True, ax=ax, **thumb_kwargs)  # type: ignore
+            _savefig('raw+roi')
 
-        if logit_cmap:
-            self.plot_with_logit_cmap(logit_cmap, show_roi=show_roi, ax=ax)
-            _savefig('custom')
-        else:
-            heatmap_kwargs = dict(
-                show_roi=show_roi,
-                interpolation=interpolation,
-                **kwargs
-            )
-            save_kwargs = dict(
-                bbox_inches='tight',
-                facecolor=ax.get_facecolor(),
-                edgecolor='none'
-            )
-            # Make heatmap plots and sliders for each outcome category
-            for i in range(self.num_classes):
-                log.info(f'Making {i+1}/{self.num_classes}...')
-                self.plot(i, heatmap_alpha=0.6, ax=ax, **heatmap_kwargs)
-                _savefig(str(i), **save_kwargs)
+            if logit_cmap:
+                self.plot_with_logit_cmap(logit_cmap, show_roi=show_roi, ax=ax)
+                _savefig('custom')
+            else:
+                heatmap_kwargs = dict(
+                    show_roi=show_roi,
+                    interpolation=interpolation,
+                    **kwargs
+                )
+                save_kwargs = dict(
+                    bbox_inches='tight',
+                    facecolor=ax.get_facecolor(),
+                    edgecolor='none'
+                )
+                # Make heatmap plots and sliders for each outcome category
+                for i in range(self.num_classes):
+                    log.info(f'Making {i+1}/{self.num_classes}...')
+                    self.plot(i, heatmap_alpha=0.6, ax=ax, **heatmap_kwargs)
+                    _savefig(str(i), **save_kwargs)
 
-                self.plot(i, heatmap_alpha=1, ax=ax, **heatmap_kwargs)
-                _savefig(f'{i}-solid', **save_kwargs)
+                    self.plot(i, heatmap_alpha=1, ax=ax, **heatmap_kwargs)
+                    _savefig(f'{i}-solid', **save_kwargs)
 
-            # Uncertainty map
-            if self.uq:
-                log.info('Making uncertainty heatmap...')
-                self.plot_uncertainty(heatmap_alpha=0.6, ax=ax, **heatmap_kwargs)
-                _savefig('UQ', **save_kwargs)
+                # Uncertainty map
+                if self.uq:
+                    log.info('Making uncertainty heatmap...')
+                    self.plot_uncertainty(heatmap_alpha=0.6, ax=ax, **heatmap_kwargs)
+                    _savefig('UQ', **save_kwargs)
 
-                self.plot_uncertainty(heatmap_alpha=1, ax=ax, **heatmap_kwargs)
-                _savefig('UQ-solid', **save_kwargs)
+                    self.plot_uncertainty(heatmap_alpha=1, ax=ax, **heatmap_kwargs)
+                    _savefig('UQ-solid', **save_kwargs)
 
-        plt.close()
+            plt.close()
         log.info(f'Saved heatmaps for [green]{self.slide.name}')
 
     def view(self):
diff --git a/slideflow/mosaic.py b/slideflow/mosaic.py
@@ -596,16 +596,17 @@ def save(self, filename: str, **kwargs: Any) -> None:
         """
         import matplotlib.pyplot as plt
 
-        self.plot(**kwargs)
-        log.info('Exporting figure...')
-        try:
-            if not os.path.exists(os.path.dirname(filename)):
-                os.makedirs(os.path.dirname(filename))
-        except FileNotFoundError:
-            pass
-        plt.savefig(filename, bbox_inches='tight')
-        log.info(f'Saved figure to [green]{filename}')
-        plt.close()
+        with sf.util.matplotlib_backend('Agg'):
+            self.plot(**kwargs)
+            log.info('Exporting figure...')
+            try:
+                if not os.path.exists(os.path.dirname(filename)):
+                    os.makedirs(os.path.dirname(filename))
+            except FileNotFoundError:
+                pass
+            plt.savefig(filename, bbox_inches='tight')
+            log.info(f'Saved figure to [green]{filename}')
+            plt.close()
 
     def save_report(self, filename: str) -> None:
         """Saves a report of which tiles (and their corresponding slide)
diff --git a/slideflow/slide/report.py b/slideflow/slide/report.py
@@ -319,17 +319,18 @@ def __init__(
                 if b is not None and b > self.bb_threshold:
                     self.warn_txt += f'{slide},{b}\n'
 
-            if np.any(n_tiles) and self.num_tiles_chart(n_tiles):
-                with tempfile.NamedTemporaryFile(suffix='.png') as temp:
-                    plt.savefig(temp.name)
-                    pdf.image(temp.name, 107, pdf.y, w=50)
-                    plt.close()
-
-            if np.any(bb) and self.blur_chart(bb):
-                with tempfile.NamedTemporaryFile(suffix='.png') as temp:
-                    plt.savefig(temp.name)
-                    pdf.image(temp.name, 155, pdf.y, w=50)
-                    plt.close()
+            with sf.util.matplotlib_backend('Agg'):
+                if np.any(n_tiles) and self.num_tiles_chart(n_tiles):
+                    with tempfile.NamedTemporaryFile(suffix='.png') as temp:
+                        plt.savefig(temp.name)
+                        pdf.image(temp.name, 107, pdf.y, w=50)
+                        plt.close()
+
+                if np.any(bb) and self.blur_chart(bb):
+                    with tempfile.NamedTemporaryFile(suffix='.png') as temp:
+                        plt.savefig(temp.name)
+                        pdf.image(temp.name, 155, pdf.y, w=50)
+                        plt.close()
 
             # Bounding box
             pdf.set_x(20)
@@ -451,11 +452,12 @@ def num_tiles_chart(self, num_tiles: np.ndarray) -> bool:
         import matplotlib.pyplot as plt
         import seaborn as sns
         if np.any(num_tiles):
-            plt.rc('font', size=14)
-            sns.histplot(num_tiles, bins=20)
-            plt.title('Number of tiles extracted')
-            plt.ylabel('Number of slides', fontsize=16, fontname='Arial')
-            plt.xlabel('Tiles extracted', fontsize=16, fontname='Arial')
+            with sf.util.matplotlib_backend('Agg'):
+                plt.rc('font', size=14)
+                sns.histplot(num_tiles, bins=20)
+                plt.title('Number of tiles extracted')
+                plt.ylabel('Number of slides', fontsize=16, fontname='Arial')
+                plt.xlabel('Tiles extracted', fontsize=16, fontname='Arial')
             return True
         else:
             return False
@@ -472,12 +474,14 @@ def blur_chart(self, blur_arr: np.ndarray) -> bool:
             with np.errstate(divide='ignore'):
                 log_b = np.log(blur_arr)
             log_b = log_b[np.isfinite(log_b)]
-            plt.rc('font', size=14)
-            sns.histplot(log_b, bins=20)
-            plt.title('Quality Control: Blur Burden'+warn_txt)
-            plt.ylabel('Count', fontsize=16, fontname='Arial')
-            plt.xlabel('log(blur burden)', fontsize=16, fontname='Arial')
-            plt.axvline(x=-3, color='r', linestyle='--')
+
+            with sf.util.matplotlib_backend('Agg'):
+                plt.rc('font', size=14)
+                sns.histplot(log_b, bins=20)
+                plt.title('Quality Control: Blur Burden'+warn_txt)
+                plt.ylabel('Count', fontsize=16, fontname='Arial')
+                plt.xlabel('log(blur burden)', fontsize=16, fontname='Arial')
+                plt.axvline(x=-3, color='r', linestyle='--')
             return True
         else:
             return False
diff --git a/slideflow/stats/metrics.py b/slideflow/stats/metrics.py
@@ -81,21 +81,23 @@ def prc_fit(self):
 
     def save_roc(self, outdir, name):
         import matplotlib.pyplot as plt
-        auroc_str = 'NA' if not self.auroc else f'{self.auroc:.2f}'
-        sf.stats.plot.roc(self.fpr, self.tpr, f'AUC = {auroc_str}')
-        full_path = join(outdir, f'{name}.png')
-        plt.savefig(full_path)
-        plt.close()
+        with sf.util.matplotlib_backend('Agg'):
+            auroc_str = 'NA' if not self.auroc else f'{self.auroc:.2f}'
+            sf.stats.plot.roc(self.fpr, self.tpr, f'AUC = {auroc_str}')
+            full_path = join(outdir, f'{name}.png')
+            plt.savefig(full_path)
+            plt.close()
         if self.neptune_run:
             self.neptune_run[f'results/graphs/{name}'].upload(full_path)
 
     def save_prc(self, outdir, name):
         import matplotlib.pyplot as plt
-        ap_str = 'NA' if not self.ap else f'{self.ap:.2f}'
-        sf.stats.plot.prc(self.precision, self.recall, label=f'AP = {ap_str}')
-        full_path = join(outdir, f'{name}.png')
-        plt.savefig(full_path)
-        plt.close()
+        with sf.util.matplotlib_backend('Agg'):
+            ap_str = 'NA' if not self.ap else f'{self.ap:.2f}'
+            sf.stats.plot.prc(self.precision, self.recall, label=f'AP = {ap_str}')
+            full_path = join(outdir, f'{name}.png')
+            plt.savefig(full_path)
+            plt.close()
         if self.neptune_run:
             self.neptune_run[f'results/graphs/{name}'].upload(full_path)
 
diff --git a/slideflow/stats/plot.py b/slideflow/stats/plot.py
@@ -1,3 +1,4 @@
+import slideflow as sf
 import os
 import warnings
 from typing import TYPE_CHECKING, List, Optional, Union
@@ -36,28 +37,29 @@ def combined_roc(
     """
     import matplotlib.pyplot as plt
 
-    plt.clf()
-    plt.title(name)
-    colors = ('b', 'g', 'r', 'c', 'm', 'y', 'k')
-    aurocs = []
-    for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
-        fpr, tpr, threshold = metrics.roc_curve(yt, yp)
-        roc_auc = metrics.auc(fpr, tpr)
-        aurocs += [roc_auc]
-        label = f'{labels[i]} (AUC: {roc_auc:.2f})'
-        plt.plot(fpr, tpr, colors[i % len(colors)], label=label)
-    plt.legend(loc='lower right')
-    plt.plot([0, 1], [0, 1], 'r--')
-    plt.xlim([0, 1])
-    plt.ylim([0, 1])
-    plt.ylabel('TPR')
-    plt.xlabel('FPR')
-    plt.savefig(os.path.join(save_dir, f'{name}.png'))
-    if neptune_run:
-        neptune_run[f'results/graphs/{name}'].upload(
-            os.path.join(save_dir, f'{name}.png')
-        )
-    plt.close()
+    with sf.util.matplotlib_backend('Agg'):
+        plt.clf()
+        plt.title(name)
+        colors = ('b', 'g', 'r', 'c', 'm', 'y', 'k')
+        aurocs = []
+        for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
+            fpr, tpr, threshold = metrics.roc_curve(yt, yp)
+            roc_auc = metrics.auc(fpr, tpr)
+            aurocs += [roc_auc]
+            label = f'{labels[i]} (AUC: {roc_auc:.2f})'
+            plt.plot(fpr, tpr, colors[i % len(colors)], label=label)
+        plt.legend(loc='lower right')
+        plt.plot([0, 1], [0, 1], 'r--')
+        plt.xlim([0, 1])
+        plt.ylim([0, 1])
+        plt.ylabel('TPR')
+        plt.xlabel('FPR')
+        plt.savefig(os.path.join(save_dir, f'{name}.png'))
+        if neptune_run:
+            neptune_run[f'results/graphs/{name}'].upload(
+                os.path.join(save_dir, f'{name}.png')
+            )
+        plt.close()
     return aurocs
 
 
@@ -176,16 +178,17 @@ def scatter(
         yp_sub = y_pred
 
     # Perform scatter for each outcome
-    for i in range(y_true.shape[1]):
-        r_squared += [metrics.r2_score(y_true[:, i], y_pred[:, i])]
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=UserWarning)
-            p = sns.jointplot(x=yt_sub[:, i], y=yp_sub[:, i], kind="reg")
-        p.set_axis_labels('y_true', 'y_pred')
-        plt.savefig(os.path.join(data_dir, f'Scatter{name}-{i}.png'))
-        if neptune_run:
-            neptune_run[f'results/graphs/Scatter{name}-{i}'].upload(
-                os.path.join(data_dir, f'Scatter{name}-{i}.png')
-            )
-        plt.close()
+    with sf.util.matplotlib_backend('Agg'):
+        for i in range(y_true.shape[1]):
+            r_squared += [metrics.r2_score(y_true[:, i], y_pred[:, i])]
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=UserWarning)
+                p = sns.jointplot(x=yt_sub[:, i], y=yp_sub[:, i], kind="reg")
+            p.set_axis_labels('y_true', 'y_pred')
+            plt.savefig(os.path.join(data_dir, f'Scatter{name}-{i}.png'))
+            if neptune_run:
+                neptune_run[f'results/graphs/Scatter{name}-{i}'].upload(
+                    os.path.join(data_dir, f'Scatter{name}-{i}.png')
+                )
+            plt.close()
     return r_squared
diff --git a/slideflow/stats/slidemap.py b/slideflow/stats/slidemap.py
@@ -984,9 +984,10 @@ def save_plot(
         """
         import matplotlib.pyplot as plt
 
-        self.plot(**kwargs)
-        plt.savefig(filename, bbox_inches='tight', dpi=dpi)
-        plt.close()
+        with sf.util.matplotlib_backend('Agg'):
+            self.plot(**kwargs)
+            plt.savefig(filename, bbox_inches='tight', dpi=dpi)
+            plt.close()
         log.info(f"Saved 2D UMAP to [green]{filename}")
 
     def save_3d(
@@ -1013,9 +1014,10 @@ def save_3d(
         """
         import matplotlib.pyplot as plt
 
-        self.plot_3d(**kwargs)
-        plt.savefig(filename, bbox_inches='tight', dpi=dpi)
-        plt.close()
+        with sf.util.matplotlib_backend('Agg'):
+            self.plot_3d(**kwargs)
+            plt.savefig(filename, bbox_inches='tight', dpi=dpi)
+            plt.close()
         log.info(f"Saved 3D UMAP to [green]{filename}")
 
     def save_coordinates(self, path: str) -> None:
diff --git a/slideflow/util/__init__.py b/slideflow/util/__init__.py
