Converting an image to black and white

Question

This code is a project that I am working on to eventually analyze some images but right now I am purely trying to read an image and convert it to black and white. The input should be either white or something else which would end up as black so it is easy to check for this.

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np

img = mpimg.imread('./Training/Image1.png')

imgplot = plt.imshow(img)

width = img.shape[1]
height = img.shape[0]

def generateArray():
    arr = []
    for y in range(img.shape[0]):
        arr.append([])
        for x in range(img.shape[1]):
            if img[y][x][0] != 1:
                arr[y].append(1)
            else:
                arr[y].append(0)
    return arr

arr = generateArray()

plt.show()

To be clear, I desire to create an array arr that contains either a 0 if the corresponding pixel in the inputted image is white or 1 if it is anything other than white.

I am somewhat of a beginner with Python but I am confused as to why this code takes so long to run (>19 seconds on my older Mac). Is there something that I am missing that would speed it up significantly?

Answer 1

Yes, there are a few things which can be done to speed up your code.

You create local variables width and height. But then you don't use these variables. Instead you use for x in range(img.shape[1]): which ends up evaluating img.shape[1] a total of height times. Consider that img.shape retrieves the shape array from img, and then accesses the [1] member of that list. Using for x in range(width): doesn't have that overhead.

Consider arr[y].append(_). Again, this code is looking up in the arr list, the [y] element, and to that retrieved list appending a value. That done once for every pixel. If you kept a handle to the list that you are appending to, you wouldn't have to look it up each time.

def generateArray():
    arr = []
    for y in range(height):
        row = []
        for x in range(width):
            if img[y][x][0] != 1:
               row.append(1)
            else:
               row.append(0)
        arr.append(row)
    return arr

The above code should be a little faster.

Appending is a time-consuming operation. If the list is stored as an array, rather than a linked-list, then appending an element will require frequent reallocation of the array capacity, and copying of all of the elements to the new storage area. We can prevent this reallocating and copying by allocating storage for all of the data all at once:

def generateArray():
    arr = [None] * height          # Correctly sized array of rows
    for y in range(height):
        row = [0] * width          # Correctly sized row, filled with 0's
        for x in range(width):
            if img[y][x][0] != 1:
               row[x] = 1
        arr[y] = row
    return arr

Here, the row array starts off with [0, 0, 0, 0, ..., 0, 0, 0] so it is only necessary to set the elements to 1 where needed.

---

Looping and list lookup can be slow. Consider:

        for x in range(width):
            if img[y][x][0] != 1:
                row[x] = 1

which loops over all the x values of a row, and for every x value, looks up the row img[y], and then the pixel in the row img[y][x], for each pixel in that row, and then accesses the [0] element of that pixel. That's a lot of lookups. If instead we used:

        for x, pixel in enumerate(img[y]):
            if pixel[0] != 1:
                row[x] = 1

we'd be iterating over the row of pixels, getting their values without the double list lookup.

We could also use list comprehension to build the list. Since each row is of a fixed, known length, the list comprehension can allocate the correctly sized list, and fill in the elements one-at-a-time, without indexing or appending.

def generateArray():
    arr = [None] * height
    for y in range(height):
        arr[y] = [ 1 if pixel[0] != 1 else 0 for pixel in img[y] ]
    return arr

That also applies to building each row.

def generateArray():
    return [ [ 1 if pixel[0] != 1 else 0 for pixel in row ] for row in img ]

---

The function generateArray() uses a global variable img. It would be better to pass img into the function.

def generateArray(img):
    # code which uses img

---

generateArray() is a terrible name for the function. Maybe generateMaskFromImage()?

---

The output arr is not being used by your code. If you omitted the line:

arr = generateArray()

your code would run much, much faster. ;-)

Answer 2

While the answer by @AJNeufeld correctly explains the shortcomings regarding normal Python in your code, there is indeed a faster way using numpy.

First, you can build a mask of all pixels that have the value 1:

mask = (img == [1, 1, 1]).all(axis=-1)

This mask is already the same as the output of your generateArray.

Executing your function takes about 720 ms ± 17.9 ms on my machine with some image file I had lying around (PNG, 471 x 698 pixels), whereas the generation of this mask is a hundred times faster at 7.37 ms ± 96.6 µs.

Now we construct two images, one completely black and one completely white:

black = np.zeros_like(img)
white = np.ones_like(img)

And then you can use numpy.where to conditionally on the mask use those two images:

img_bw = np.where(mask[:,:,None], white, black)

The weird looking slicing with None just makes sure that the mask has the right dimensionality (i.e. three values per pixel, instead of just one).

Putting it all together:

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np

def to_bw(img):
    """Return a black and white copy of `img`.

    Pixels which are white in the original image stay white,
    all other pixels are turned black.
    """
    mask = (img == [1, 1, 1]).all(axis=-1)
    black = np.zeros_like(img)
    white = np.ones_like(img)
    return np.where(mask[:,:,None], white, black)

if __name__ == "__main__":
    img = mpimg.imread('./Training/Image1.png')
    img_bw = to_bw(img)
    plt.imshow(img_bw)
    plt.show()

The whole to_bw function only takes 14 ms ± 324 µs.