使用win32gui进行屏幕捕获时出现内存泄漏求解决方案

以下代码使用win32系列API实现桌面截屏，但运行时内存持续增长，推测存在内存泄漏，排查后未发现明显问题，寻求解决方案。运行环境：Python 3.10.11、pywin32、Windows 10。

import numpy as np
import win32gui, win32ui, win32con
import time 
import cv2
import psutil

def memory_usage(message: str = 'debug'):
    # current process RAM usage
    p = psutil.Process()
    rss = p.memory_info().rss / 2 ** 20 # Bytes to MB
    print(f"[{message}] memory usage: {rss: 10.5f} MB")

class ImageSearch:


    # constructor
    def __init__(self, windowx=0, windowy=0, windowr=0):
        self.hwnd = win32gui.GetDesktopWindow()  # Get the desktop window handle

    def screencapture(self):
        # Get the window size
        window_rect = win32gui.GetWindowRect(self.hwnd)
        self.left = window_rect[0]
        self.top = window_rect[1]
        self.weight = window_rect[2] - window_rect[0]
        self.height = window_rect[3] - window_rect[1]

        # Get the window image data
        wDC = win32gui.GetWindowDC(self.hwnd)
        dcObj = win32ui.CreateDCFromHandle(wDC)
        cDC = dcObj.CreateCompatibleDC()
        dataBitMap = win32ui.CreateBitmap()
        dataBitMap.CreateCompatibleBitmap(dcObj, self.weight, self.height)
        cDC.SelectObject(dataBitMap)
        cDC.BitBlt((0, 0), (self.weight, self.height), dcObj, (self.left, self.top), win32con.SRCCOPY)

        # Convert the raw data into a format OpenCV can read
        signedIntsArray = dataBitMap.GetBitmapBits(True)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.weight, 4)

        # Clean up resources
        dcObj.DeleteDC()
        cDC.DeleteDC()
        win32gui.ReleaseDC(self.hwnd, wDC)
        win32gui.DeleteObject(dataBitMap.GetHandle())

        img = img[..., :3]
        img = np.ascontiguousarray(img)

        return img

a = ImageSearch()
while True:
    memory_usage('#1')
    abc = a.screencapture()

问题原因 #

内存泄漏的核心问题出在dataBitMap.GetBitmapBits(True)调用上：当传入True参数时，pywin32会返回一个Python字节对象，但该对象关联的原生内存块没有被正确标记为可回收，导致每次循环都会残留这部分内存，最终表现为内存持续增长。另外原代码中还有个小错误：self.weight是拼写错误，应为self.width，虽不影响泄漏，但建议修正。

解决方案 #

方案一：修复原生Win32API的内存管理

修改GetBitmapBits的调用方式，手动管理内存句柄，避免泄漏：

def screencapture(self):
    # 获取窗口尺寸
    window_rect = win32gui.GetWindowRect(self.hwnd)
    self.left = window_rect[0]
    self.top = window_rect[1]
    self.width = window_rect[2] - window_rect[0]  # 修正拼写错误
    self.height = window_rect[3] - window_rect[1]

    # 获取窗口图像数据
    wDC = win32gui.GetWindowDC(self.hwnd)
    dcObj = win32ui.CreateDCFromHandle(wDC)
    cDC = dcObj.CreateCompatibleDC()
    dataBitMap = win32ui.CreateBitmap()
    dataBitMap.CreateCompatibleBitmap(dcObj, self.width, self.height)
    # 保存DC原有位图，后续恢复避免泄漏
    old_bitmap = cDC.SelectObject(dataBitMap)
    cDC.BitBlt((0, 0), (self.width, self.height), dcObj, (self.left, self.top), win32con.SRCCOPY)

    # 手动处理位图内存，解决泄漏
    import ctypes
    buf_size = self.width * self.height * 4
    # 获取内存句柄而非直接返回字节对象
    h_bits = dataBitMap.GetBitmapBits(False)
    try:
        signedIntsArray = ctypes.c_buffer(buf_size)
        ctypes.memmove(signedIntsArray, h_bits, buf_size)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.width, 4)
    finally:
        # 手动释放内存句柄
        win32gui.DeleteObject(h_bits)

    # 恢复DC原有位图
    cDC.SelectObject(old_bitmap)

    # 清理资源
    dcObj.DeleteDC()
    cDC.DeleteDC()
    win32gui.ReleaseDC(self.hwnd, wDC)
    win32gui.DeleteObject(dataBitMap.GetHandle())

    img = img[..., :3]
    img = np.ascontiguousarray(img)

    return img

方案二：改用成熟的截屏库

如果不想手动处理底层内存管理，可以用mss库替代原生Win32API，它在内存优化上做了完善处理，代码更简洁：

import numpy as np
import cv2
import psutil
import mss

def memory_usage(message: str = 'debug'):
    p = psutil.Process()
    rss = p.memory_info().rss / 2 ** 20
    print(f"[{message}] memory usage: {rss: 10.5f} MB")

class ImageSearch:
    def __init__(self):
        self.sct = mss.mss()

    def screencapture(self):
        # 捕获主显示器（monitors[1]为主屏，monitors[0]为全屏幕）
        monitor = self.sct.monitors[1]
        img = np.array(self.sct.grab(monitor))
        img = img[..., :3]  # 移除alpha通道
        img = np.ascontiguousarray(img)
        return img

a = ImageSearch()
while True:
    memory_usage('#1')
    abc = a.screencapture()

安装mss：pip install mss

内容的提问来源于stack exchange，提问作者mahaha3211