CODEWITHSUNDEEP
pythonoptimizationnumpyscipyinterpolation
Andrew
Andrew

Reputation: 2892

Speed up this interpolation in python

I have an image processing problem I'm currently solving in python, using numpy and scipy. Briefly, I have an image that I want to apply many local contractions to. My prototype code is working, and the final images look great. However, processing time has become a serious bottleneck in our application. Can you help me speed up my image processing code?

I've tried to boil down our code to the 'cartoon' version below. Profiling suggests that I'm spending most of my time on interpolation. Are there obvious ways to speed up execution?

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)

Vague explanation of what the code does:

We start with a pixellated 2D image, and a set of arbitrary (x, y) points in our image that don't generally fall on an integer grid. For each (x, y) point, I want to multiply the image by a small mask centered precisely on that point. Next we contract/expand the masked region by a finite amount, before finally adding this processed sub-image to a final image, which may not have the same pixel size as the original image. (Not my finest explanation. Ah well).

Upvotes: 1

Views: 1728

Answers (1)

matehat
matehat

Reputation: 5374

I'm pretty sure that, as you said, the bulk of the calculation time happens in interpolate.map_coordinates(…), which gets called once for every iteration on center_points, here 10,000 times. Generally, working with the numpy/scipy stack, you want the repetitive task over a large array to happen in native Numpy/Scipy functions -- i.e. in a C loop over homogeneous data -- as opposed to explicitely in Python.

One strategy that might speed up the interpolation, but that will also increase the amount of memory used, is :

  • First, fetch all the subimages (here named masked_image) in a 3-dimensional array (window_radius x window_radius x center_points.size)
  • Make a ufunc (read that, it's useful) that wraps the work that has to be done on each subimage, using numpy.frompyfunc, which should return another 3-dimensional array (subgrid_radius[0] x subgrid_radius[1] x center_points.size). In short, this creates a vectorized version of the python function, that can be broadcast element-wise on an array.
  • Build the final image by summing over the third dimension.

Hope that gets you closer to your goals!

Upvotes: 3

Related Questions