pcolormesh内存使用情况

Question

我最近一直在为这个问题苦恼。我有一组图片，我对这些图片的X和Y坐标进行了一些数学运算，然后用pcolormesh来绘制新的图片。所有的计算我都已经完成了，现在只需要加载新的X和Y坐标，并在pcolormesh中使用图片的颜色。

这些图片的大小是2048x2448像素（大约5百万像素），第一张图片处理得很快，但之后的每一张图片，脚本的运行速度就变慢，而且占用的内存也越来越多。我尝试了一些垃圾回收的方法，但效果不太好。

我的脚本：

import numpy as np
from PIL import Image
import cPickle as pickle
import matplotlib.pyplot as plt
import os

# TRY forced garbage collection!
import gc

def rectify(cam_files, cam_rec_files, rec_file_location):
    ''' cam_files is a dictionary that contains the filenames with the camera-names as index
            example: {'KDXX04C' : C:\Users\Yorian\Desktop\TU\Stage Shore\python_files\Rectify, metadata and upload\v3\archive\KDXXXXX\original\snap\1381383000\{filename}.jpg }

        cam_rec_files_dir contains a dictionary, cameraName : fileLocation
            example: {'KDXX04C' : C:\Users\Yorian\Desktop\TU\Stage Shore\python_files\Rectify, metadata and upload\v3\camdata\KDXXXXX\KDXX04C.pkl }

        rec_file_location is a string that shows where the new rectification needs to be saved '''
    fig, ax = plt.subplots(1, 1, figsize=(60,90))

    for camname in cam_files:
        img = Image.open(cam_files[camname])
        img = np.asarray(img, dtype=np.uint8)
        height, width, channels = img.shape

        # load plot data
        fh = open(cam_rec_files[camname], 'rb')
        X = pickle.load(fh)
        Y = pickle.load(fh)

        masks = [X<=0, X>1500, Y>4000, Y<-4000]
        total_mask = masks[0] | masks[1] | masks[2] | masks[3]
        first_false = np.argwhere(total_mask == 0)
        start = int(first_false[0]/width)

        rgb = img.reshape((-1,3))/255.0
        rgba = np.concatenate((rgb, np.ones((rgb.shape[0],1), dtype=np.uint8)), axis=1)
        rgba[total_mask,3] = 0
        rgba = rgba.reshape((height,width,4))[:,:-1,:]
        rgba = rgba.reshape((-1,4))

        plotimg = ax.pcolormesh(X.reshape((height, width))[start:,:], Y.reshape((height, width))[start:,:], img.mean(-1)[start:,:], cmap='Greys') # img.mean(-1)

        plotimg.set_array(None)
        plotimg.set_edgecolor('none')
        plotimg.set_facecolor(rgba[(start*(width-1)):,:])

        fh.close()

    plt.savefig(rec_file_location)
    gc.collect()

这个脚本在处理六张图片时还可以，但当我尝试处理八张图片时，就出现了内存不足的问题（我使用的是64位的Python，电脑上有12GB的内存，我原本以为这已经足够了）。

有没有人知道怎么解决这个问题？

Answer 1

如果你清空ax中存储的数据，就可以减少所需的内存。你可以在Joe Kington的rectify函数最后加上ax.clear()。这样，程序只需要一个pcolormesh的内存。

Answer 2

简单来说，如果你在使用 pyplot 接口并且想要生成很多图形但不想显示它们，当你用完某个图形后，就调用 plt.close(fig)。

每次你调用 rectify 函数时，都会生成一个新的（非常大的！！）图形，并且这个图形会被 保存在内存中。pyplot 会保留这个图形的引用，以便你在调用 plt.show() 时可以显示它。你可以选择调用 plt.close(fig)，或者在不使用 pyplot 状态机的情况下创建图形。（fig.clf() 也可以，但会保留空图形的引用。）

另外，考虑到你是在读取图像文件，你的值应该是在一个规则的 x 和 y 网格上。如果是这样，使用 imshow 会比 pcolormesh 更快且更节省内存。

---

举个例子，关于第一个问题，你的 rectify 函数基本上是这样做的，而你可能会重复调用它（就像下面的循环那样）：

import numpy as np
import matplotlib.pyplot as plt

def rectify():
    fig, ax = plt.subplots()
    ax.pcolormesh(np.random.random((10,10)))
    fig.savefig('blah.png')

for _ in range(10):
    rectify()

plt.show()

注意，我们会看到10个图形弹出来。pyplot 会保留这个图形的引用，以便可以用 show 显示。

如果你想从 pyplot 的状态机中移除这个图形，可以调用 plt.close(fig)。

例如，如果你这样做，就不会显示任何图形：（每个图形在你通过调用 plt.close(fig) 从 pyplot 的图形管理器中移除后，会像你预期的那样被垃圾回收。）

import numpy as np
import matplotlib.pyplot as plt

def rectify():
    fig, ax = plt.subplots()
    ax.pcolormesh(np.random.random((10,10)))
    fig.savefig('blah.png')
    plt.close(fig)

for _ in range(10):
    rectify()

plt.show()

另外，你也可以绕过 pyplot，直接创建图形和画布。这样就不会涉及到 pyplot 的图形管理器，图形实例会像你预期的那样被垃圾回收。不过，这种方法比较繁琐，并且假设你对 matplotlib 的幕后工作有一些了解：

import numpy as np
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

# Don't need this, but just to demonstrate that `show()` has no effect...
import matplotlib.pyplot as plt

def rectify():
    fig = Figure()
    FigureCanvas(fig)
    ax = fig.add_subplot(1,1,1)
    ax.pcolormesh(np.random.random((10,10)))
    fig.savefig('blah.png')

for _ in range(10):
    rectify()

plt.show()