Matt
Matt

Reputation: 75

Most efficient way to stretch an image in Python NumPy

I want a function to take in an image as a numpy array and remap the values to a new range (0, 1) based on a specified maximum and minimum value from the input range. I've got a working function, but I'm iterating through the array and it takes about 10 seconds to complete. Is there a more efficient way to perform this task? Maybe some built in numpy function that I'm not aware of?

This is what I've got:

import numpy as np

def stretch(image, minimum, maximum):
    dY = image.shape[0]
    dX = image.shape[1]
    r = maximum - minimum
    image = image.flatten()
    for i in range(image.size):
        if image[i] > maximum or image[i] < minimum:
            image[i] = 1. or 0.
        else:
            image[i] = (image[i] - minimum) / r
    return image.reshape(dY, dX)

I've also tried a version of the above using numpy.nditer instead of manually iterating with the for loop but that seems to be about four times as slow (~40 seconds).

Is there a more efficient way to do this that I'm overlooking? The images I'm working with are about 16 MP. (3520, 4656)

Upvotes: 1

Views: 6959

Answers (1)

Alex
Alex

Reputation: 19104

Bug #1

You have a bug in your code.

image[i] = 1. or 0. always evaluates to 1.0 because 1. acts truthy.

Instead that block should look like:

if image[i] < minimum:
    image[i] = 0.
elif image[i] > maximum:
    image[i] = 1
else:
    image[i] = (image[i] - minimum) / r

Bug #2

If your original array is of dtype=int and you place values into it they will be coerced to ints. This means that any floats will be rounded down.

a = np.array([1])
a[0] = 0.5
a

returns

array([0])

This can be addressed using the vectorized solution below.

Solution

In general try not to use loops when dealing with NumPy arrays. Using vectorized functions can be much faster and more readable.

def stretch(image, minimum, maximum):
    image = (image - minimum) / (maximum - minimum)
    image[image < 0] = 0
    image[image > 1] = 1
    return image

An example (updated to int which is more realistic for image as @MrT points out):

a = np.arange(1, 4).reshape(2, 2)
stretch(a, 1, 3)

returns

array([[0. , 0.5],
       [1. , 1. ]])

Upvotes: 2

Related Questions