使用颜色条删除图像不会释放 matplotlib 中的内存？

Posted 2023-02-23

【中文标题】使用颜色条删除图像不会释放 matplotlib 中的内存？

【英文标题】：remove image with colorbar won't release the memory in matplotlib?

【发布时间】：2022-01-24 04:01:43

【问题描述】：

我正在尝试从图形中删除图像并释放内存。图片不加colorbar时，可以成功释放内存，如果加colorbar则失败。在下面的演示代码中：

单击按钮添加颜色条将为图中的一个图像添加一个颜色条。 单击按钮删除将从图中删除一个图像（和相关的颜色条）。

每次删除图片，相关的colorbar也被删除了，所以不知道为什么内存回收失败，我猜想添加colorbar的时候一定是对图片有一些额外的引用，失败了内存回收。

import numpy as np
from PyQt5 import QtWidgets
from memory_profiler import profile
import matplotlib
from matplotlib.figure import Figure
import matplotlib.cm as cm
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from matplotlib.axes._axes import Axes

matplotlib.use("Qt5Agg")


class MplCanvas(FigureCanvasQTAgg):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.axe = self.fig.add_subplot(1, 1, 1, label='good')
        super().__init__(self.fig)


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()

        layout = QtWidgets.QVBoxLayout()

        self.canvas = MplCanvas(self, width=5, height=4, dpi=100)
        self.axe = self.canvas.axe

        layout.addWidget(self.canvas)

        self.pushButton_addColorBar = QtWidgets.QPushButton('Add ColorBar')
        layout.addWidget(self.pushButton_addColorBar)
        self.pushButton_remove = QtWidgets.QPushButton('remove')
        layout.addWidget(self.pushButton_remove)

        widget = QtWidgets.QWidget()
        widget.setLayout(layout)
        self.setCentralWidget(widget)

        self.pushButton_remove.clicked.connect(self.removeImage)
        self.pushButton_addColorBar.clicked.connect(self.createColorBar)

        self.pcolormesh_test()

    def pcolormesh_test(self):
        """add two images"""
        delta = 0.01
        x = y = np.arange(-3.0, 3.0, delta)
        X, Y = np.meshgrid(x, y)
        Z1 = np.exp(-X ** 2 - Y ** 2)
        Z2 = np.exp(-(X - 1) ** 2 - (Y - 1) ** 2)
        Z = (Z1 - Z2) * 2

        im = self.axe.pcolormesh(X, Y, Z, cmap=cm.viridis, shading='auto')
        im.set_clim(vmax=np.amax(Z), vmin=np.amin(Z))

        Zx = (Z1 + Z2) * 2
        imx = self.axe.pcolormesh(X, Y, Zx, cmap=cm.Blues, shading='auto')
        imx.set_clim(vmax=np.amax(Zx), vmin=np.amin(Zx))

    def createColorBar(self):
        """ to create a color bar for an image. """
        axe = self.axe
        fig = axe.get_figure()
        images = self.getImages(axe)

        for image in images:
            if not image.colorbar:  # color bar doesn't exist
                inset_axe = axe.inset_axes([1.0, 0, 0.05, 1], transform=axe.transAxes)
                fig.colorbar(image, ax=axe, cax=inset_axe)
                break  # each trigger create one colorbar for one image

        self.reDraw()

    @profile
    def removeImage(self, checked):
        """
        Usage:
            * each trigger remove one image

        """
        images = self.getImages(self.axe)
        # print(f'images=images')

        if images:
            image = images[-1]
            color_bar = image.colorbar
            if color_bar:
                color_bar.remove()
                del color_bar

            # remove image
            image.remove()
            del image

        self.reDraw()

    def getImages(self, axe: Axes):
        """to obtain the image list in the axe"""
        images = []
        images.extend(axe.images)
        images.extend(axe.collections)

        return images

    def reDraw(self):
        self.canvas.draw_idle()
        self.canvas.flush_events()


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)

    window = MainWindow()
    window.show()
    sys.exit(app.exec_())

【参考方案1】：

我已经找到解决方案，并发布答案以寻求帮助。

我们需要在removeImage()方法的末尾添加gc.collect()。那么当图像被删除时，内存可以被回收。