Ускорьте эту интерполяцию в Python - PullRequest
1 голос
/ 17 ноября 2011

У меня есть проблема с обработкой изображений, которую я сейчас решаю на python, используя numpy и scipy. Вкратце, у меня есть изображение, к которому я хочу применить множество локальных сокращений. Мой прототип кода работает, и окончательные изображения выглядят великолепно. Однако время обработки стало серьезным узким местом в нашем приложении. Можете ли вы помочь мне ускорить мой код обработки изображений?

Я пытался свести наш код к версии мультфильма ниже. Профилирование предполагает, что я трачу большую часть своего времени на интерполяцию. Есть ли очевидные способы ускорить выполнение?

import cProfile, pstats
import numpy
from scipy.ndimage import interpolation

def get_centered_subimage(
    center_point, window_size, image):
    x, y = numpy.round(center_point).astype(int)
    xSl = slice(max(x-window_size-1, 0), x+window_size+2)
    ySl = slice(max(y-window_size-1, 0), y+window_size+2)
    subimage = image[xSl, ySl]
    interpolation.shift(
        subimage, shift=(x, y)-center_point, output=subimage)
    return subimage[1:-1, 1:-1]

"""In real life, this is experimental data"""
im = numpy.zeros((1000, 1000), dtype=float)
"""In real life, this mask is a non-zero pattern"""
window_radius = 10
mask = numpy.zeros((2*window_radius+1, 2*window_radius+1), dtype=float)
"""The x, y coordinates in the output image"""
new_grid_x = numpy.linspace(0, im.shape[0]-1, 2*im.shape[0])
new_grid_y = numpy.linspace(0, im.shape[1]-1, 2*im.shape[1])


"""The grid we'll end up interpolating onto"""
grid_step_x = new_grid_x[1] - new_grid_x[0]
grid_step_y = new_grid_y[1] - new_grid_y[0]
subgrid_radius = numpy.floor(
    (-1 + window_radius * 0.5 / grid_step_x,
     -1 + window_radius * 0.5 / grid_step_y))
subgrid = (
    window_radius + 2 * grid_step_x * numpy.arange(
        -subgrid_radius[0], subgrid_radius[0] + 1),
    window_radius + 2 * grid_step_y * numpy.arange(
        -subgrid_radius[1], subgrid_radius[1] + 1))
subgrid_points = ((2*subgrid_radius[0] + 1) *
                  (2*subgrid_radius[1] + 1))

"""The coordinates of the set of spots we we want to contract. In real
life, this set is non-random:"""
numpy.random.seed(0)
num_points = 10000
center_points = numpy.random.random(2*num_points).reshape(num_points, 2)
center_points[:, 0] *= im.shape[0]
center_points[:, 1] *= im.shape[1]

"""The output image"""
final_image = numpy.zeros(
    (new_grid_x.shape[0], new_grid_y.shape[0]), dtype=numpy.float)

def profile_me():
    for m, cp in enumerate(center_points):
        """Take an image centered on each illumination point"""
        spot_image = get_centered_subimage(
            center_point=cp, window_size=window_radius, image=im)
        if spot_image.shape != (2*window_radius+1, 2*window_radius+1):
            continue #Skip to the next spot
        """Mask the image"""
        masked_image = mask * spot_image
        """Resample the image"""
        nearest_grid_index = numpy.round(
                (cp - (new_grid_x[0], new_grid_y[0])) /
                (grid_step_x, grid_step_y))
        nearest_grid_point = (
            (new_grid_x[0], new_grid_y[0]) +
            (grid_step_x, grid_step_y) * nearest_grid_index)
        new_coordinates = numpy.meshgrid(
            subgrid[0] + 2 * (nearest_grid_point[0] - cp[0]),
            subgrid[1] + 2 * (nearest_grid_point[1] - cp[1]))
        resampled_image = interpolation.map_coordinates(
            masked_image,
            (new_coordinates[0].reshape(subgrid_points),
             new_coordinates[1].reshape(subgrid_points))
            ).reshape(2*subgrid_radius[1]+1,
                      2*subgrid_radius[0]+1).T
        """Add the recentered image back to the scan grid"""
        final_image[
            nearest_grid_index[0]-subgrid_radius[0]:
            nearest_grid_index[0]+subgrid_radius[0]+1,
            nearest_grid_index[1]-subgrid_radius[1]:
            nearest_grid_index[1]+subgrid_radius[1]+1,
            ] += resampled_image

cProfile.run('profile_me()', 'profile_results')
p = pstats.Stats('profile_results')
p.strip_dirs().sort_stats('cumulative').print_stats(10)

Неясное объяснение того, что делает код:

Мы начинаем с пиксельного 2D-изображения и набора произвольных (x, y) точек в нашем изображении, которые обычно не попадают в целочисленную сетку. Для каждой (x, y) точки я хочу умножить изображение на маленькую маску с центром точно в этой точке. Затем мы сокращаем / расширяем замаскированную область на конечную величину, прежде чем окончательно добавить это обработанное подизображение к конечному изображению, которое может не иметь такой же размер пикселя, как исходное изображение. (Не мое лучшее объяснение. Ах, хорошо).

1 Ответ

3 голосов
/ 17 ноября 2011

Я почти уверен, что, как вы сказали, основная часть времени вычислений происходит в interpolate.map_coordinates(…), который вызывается один раз для каждой итерации в center_points, здесь 10000 раз.Как правило, работая со стеком Numpy / Scipy, вы хотите, чтобы повторяющаяся задача для большого массива выполнялась в собственных функциях Numpy / Scipy, т. Е. В цикле C над однородными данными, а не явно в Python.

Одна из стратегий, которая может ускорить интерполяцию, но также увеличит объем используемой памяти:

  • Сначала выберите все подизображения (здесь они называются masked_image) в 3-размерный массив (window_radius x window_radius x center_points.size)
  • Создайте ufunc (прочитайте это, это полезно), который оборачивает работу, которую необходимо выполнить для каждого подизображения, используя numpy.frompyfunc , который должен вернуть другой трехмерный массив (subgrid_radius[0] x subgrid_radius[1] x center_points.size).Короче говоря, это создает векторизованную версию функции python, которую можно транслировать поэлементно в массиве.
  • Создание окончательного изображения путем суммирования по третьему измерению.

Надеюсь, это приблизит вас к вашим целям!

...