python - 图像 block 的多重处理

Question

我有一个函数，它必须遍历图像的各个像素并计算一些几何图形。这个函数需要很长时间才能运行(在 24 兆像素的图像上大约需要 5 小时)，但看起来应该很容易在多核上并行运行。但是，我终其一生都无法找到一个使用 Multiprocessing 包做这样的事情的有据可查、解释清楚的例子。这是我现在作为玩具示例运行的代码:

import numpy as np
import matplotlib.pyplot as plt
from scipy import misc
from skimage import color
import multiprocessing 
from multiprocessing import Process

#Some dumb stand in function for this exercise
def dumb_func(image):
    ny, nx = image.shape
    temp = np.empty_like(image)

    for y in range(ny):
        for x in range(nx):
            temp[y, x] = np.square(image[y, x])

    return temp

#Convert image to greyscale
img = color.rgb2gray(misc.ascent())

#Resize the image
ns = 2048 #Pixel size
img = misc.imresize(img, size = (ns, ns))


#Split the image into equal chunks...not sure how this works for arrays that
#are weird shapes and aren't the same size in each dimension

divs = 4
init_split = np.array_split(img, divs, axis = 0)
side = init_split[0].shape[0]
chunked = np.empty((divs, divs, side, side))
cur = 0
for i in range(divs):
    split = np.array_split(init_split[i], divs, axis = 1)
    for j in range(divs):
        chunked[i, j, :, :] = split[j]
        cur +=1

#Pull core count and divide by two to be safe
cores = int(multiprocessing.cpu_count() / 2)

result = np.empty_like(chunked)
idxs = np.array(np.meshgrid(np.arange(0, divs, 1), 
                            np.arange(0, divs, 1))).T.reshape(-1, 2)

基本上这段代码会加载图像，将其转换为灰度，使其变大，然后将其分块。分块数组的形状为 (i, j, ny, nx)，其中 i 和 j 是标识我正在处理的图像 block 的索引，而 ny,nx 描述每个 block 的像素大小。

此外，我正在创建一个名为 idxs 的数组，它将所有可能的索引存储到分块数组中，以提取分块图像。

我想做的是在 block 上并行运行一个函数(在本例中以 dumb_func 为例)，并将结果存储在相同形状的结果数组中。我想象的这样做的方式是遍历 idxs 数组并分配进程属于这些索引的 block 直到核心数，等待这些核心完成，然后为核心提供更多进程直到完成。我被卡住了，因为我无法 A) 弄清楚如何访问函数中的返回值，以及 B) 如何处理我可能有 16 个 block 和 5 个核心导致最后一次迭代只需要一个进程的情况。

我该怎么做呢？在过去的 6-7 个小时里，我阅读了有关多处理池、进程、映射、星图等的内容……但我终其一生都无法理解如何实现它。

为 Reedininer 编辑:

这是我更新后的代码，运行没有错误。但是 new_data 数组永远不会更新。我用 100 的值填充它，在例程结束时，new_data 正是它的初始化方式。

import numpy as np
import matplotlib.pyplot as plt
from scipy import misc
from multiprocessing import Process, JoinableQueue
from time import time

#SOme dumb stand in function for this exercise
def dumb_func(q, new_data):
    while True:
        index, image = q.get()
        temp = image **2

        new_data[index[0], index[1], :, :] = temp
        q.task_done()

if __name__ == "__main__":
    start = time()
    q = JoinableQueue()
    img = misc.ascent()
    #Resize the image
    ns = 2048 #Pixel size
    img = misc.imresize(img, size = (ns, ns))
    #Split the image into equal chunks...not sure how this works for arrays that
    #are weird shapes and aren't the same size in each dimension

    divs = 4
    init_split = np.array_split(img, divs, axis = 0)
    side = init_split[0].shape[0]
    chunked = np.empty((divs, divs, side, side))
    cur = 0
    for i in range(divs):
        split = np.array_split(init_split[i], divs, axis = 1)
        for j in range(divs):
            chunked[i, j, :, :] = split[j]
            cur +=1

    new_data = np.full(chunked.shape, 100)
    idxs = np.array(np.meshgrid(np.arange(0, divs, 1), 
                                np.arange(0, divs, 1))).T.reshape(-1, 2)

    for i in range(len(idxs)):
        q.put((idxs[i], chunked[idxs[i][0], idxs[i][1], :, :]))

    print ('starting workers')

    worker_count = len(idxs)
    processes = []
    for i in range(worker_count):
        p = Process(target=dumb_func, args=[q, new_data])
        p.daemon = True
        p.start()
    print('main thread waiting')
    q.join()

    end = time()
    print('{:.3f} seconds elapsed'.format(end - start))

Answer 1

我会做这样的事情，从依赖关系开始:

from multiprocessing import Pool
import numpy as np
from PIL import Image

# and some for testing
from random import random
from time import sleep

首先我定义了一个函数来将图像分成“ block ”，有点像你所说的:

def chunkit(ys, xs, blocksize=64):
    for y in range(0, ys, blocksize):
        yt = (y, min(ys, y + blocksize))
        for x in range(0, xs, blocksize):
            xt = (x, min(xs, x + blocksize))
            yield yt, xt

这是一个惰性迭代器，所以这可以持续一段时间。

然后我定义我的辅助函数:

def dumb_func(cc):
    (y0,y1), (x0,x1) = cc
    # convert to floats for ease of processing
    chunk = image[y0:y1,x0:x1] / 255.
    # random slow down for testing
    # sleep(random() ** 6)
    res = chunk ** 2
    # convert back to bytes for efficiency
    return cc, (res * 255).astype(np.uint8)

我确保源数组尽可能接近原始格式以提高效率，并以相同的格式将其发回(如果您显然要处理其他像素格式，这可能需要一些调整)。

然后我把它放在一起:

if __name__ == '__main__':
    source = Image.open('tmp.jpeg')
    image = np.asarray(source)
    print("loaded", image.shape, image.dtype)

    with Pool() as pool:
        resit = pool.imap_unordered(
            dumb_func, chunkit(*image.shape[:2]))

        output = np.empty_like(image)
        for cc, res in resit:
            (y0,y1), (x0,x1) = cc
            output[y0:y1,x0:x1] = res

    im = Image.fromarray(output, 'RGB')
    im.save('out.jpeg')

这会在几秒钟内完成一张 15Mpixel 的图像，其中大部分时间用于加载/保存图像。数组跨度和缓存友好性可能会更聪明，但希望对您有所帮助!

注意:我认为这段代码依赖于 CPython Unix 风格的进程 fork 语义来确保图像在进程之间有效共享。不确定如果你在其他东西上运行它会发生什么