我的图像处理代码的瓶颈在哪里？

Question

我编写这个脚本是为了对大量 PNG 文件（总共大约 1500 个）进行一些图像处理。它们被组织成子目录。

这是我的代码：

from PIL import Image
import os

path = "/Some/given/path"

file_list = []
counter = 1

for root, dirs, files in os.walk(path):
    for file in files:
        if file.endswith(".png"):
            temp_file = {"path": os.path.join(root, file), "name": file}
            file_list.append(temp_file)

for curr_file in file_list:
    img = Image.open(curr_file["path"])
    img = img.convert("RGBA")
    val = list(img.getdata())
    new_data = []
    for item in val:
        if item[3] == 0:
            new_data.append(item)
        else:
            new_data.append((0, 0, 0, 255))
        img.putdata(new_data)
    file_name = "transform" + str(counter) + ".png"
    replaced_text = curr_file["name"]
    new_file_name = curr_file["path"].replace(replaced_text, file_name)
    img.save(new_file_name)
    counter += 1

文件夹结构如下：

Source folder
     -- folder__1
        -- image_1.png
        -- image_2.png
        -- image_3.png
     -- folder__2
        -- image_3.png
        -- image_5.png
     -- folder__3
        -- image_6.png

在对单个图像进行测试时，图像处理只需几秒钟。但是，在运行脚本时，处理 15 张图像大约需要一个小时。关于我在哪里搞砸的任何建议？

Answer 1

主要问题在这里：

new_data = []
for item in val:
    if item[3] == 0:
        new_data.append(item)
    else:
        new_data.append((0, 0, 0, 255))
    img.putdata(new_data)                   # <--

如果您要收集完整的new_data，则无需为每个像素更新img 的内容。因此，只需将该行移出循环即可：

new_data = []
for item in val:
    if item[3] == 0:
        new_data.append(item)
    else:
        new_data.append((0, 0, 0, 255))
img.putdata(new_data)                       # <--

现在，通过使用 NumPy 及其矢量化功能，彻底摆脱对所有像素的迭代：

from PIL import Image
import os
import numpy as np                          # <--

path = "/Some/given/path"

file_list = []
counter = 1

for root, dirs, files in os.walk(path):
    for file in files:
        if file.endswith(".png"):
            temp_file = {"path": os.path.join(root, file), "name": file}
            file_list.append(temp_file)

for curr_file in file_list:
    img = Image.open(curr_file["path"])
    img = img.convert("RGBA")
    img = np.array(img)                     # <--
    img[img[..., 3] != 0] = (0, 0, 0, 255)  # <--
    img = Image.fromarray(img)              # <--
    file_name = "transform" + str(counter) + ".png"
    replaced_text = curr_file["name"]
    new_file_name = curr_file["path"].replace(replaced_text, file_name)
    img.save(new_file_name)
    counter += 1

基本上，您将所有 alpha 通道不等于 0 的像素设置为 (0, 0, 0, 255)。这就是您在那里看到的 NumPy 单线。前后行仅用于从 Pillow Image 转换为 NumPy 数组，反之亦然。

---

编辑：如果您不想在代码中使用 NumPy，您还可以使用 Pillow 的 point 函数摆脱循环，参见。 this tutorial:

from PIL import Image
import os

path = "/Some/given/path"

file_list = []
counter = 1

for root, dirs, files in os.walk(path):
    for file in files:
        if file.endswith(".png"):
            temp_file = {"path": os.path.join(root, file), "name": file}
            file_list.append(temp_file)

for curr_file in file_list:
    img = Image.open(curr_file["path"])
    img = img.convert("RGBA")
    source = img.split()                                                # <--
    mask = source[3].point(lambda i: i > 0 and 255)                     # <--
    img.paste(Image.new("RGBA", img.size, (0, 0, 0, 255)), None, mask)  # <--
    file_name = "transform" + str(counter) + ".png"
    replaced_text = curr_file["name"]
    new_file_name = curr_file["path"].replace(replaced_text, file_name)
    img.save(new_file_name)
    counter += 1

----------------------------------------
System information
----------------------------------------
Platform:      Windows-10-10.0.16299-SP0
Python:        3.9.1
NumPy:         1.20.2
Pillow:        8.1.2
----------------------------------------

Answer 2

您可以使用 snakeviz 库来分析您的代码 -

Snakeviz - https://jiffyclub.github.io/snakeviz/

python -m cProfile -o program.prof my_program.py

生成配置文件后，您可以可视化并查看哪个功能/哪条线需要更多时间。

snakeviz program.prof