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

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

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

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

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

配置文件生成后，您可以可视化并查看哪条 function/which 行花费了更多时间。

snakeviz program.prof

Answer 2

主要问题位于此处：

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)                   # <--

您不需要为每个像素更新 img 的内容，如果您收集的是完整的 new_data。所以，只需将该行移到循环外：

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 function, cf. 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
----------------------------------------

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

Where is the bottleneck in my image manipulation code?

python

for-loop

image-processing

python-imaging-library

python-3.x