Make an image where pixels are colored if they are prime

Question

This code creates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't. It works, but I wonder if it is well written or if it could be made faster.

Example:

Code:

import math

from PIL import Image


def is_prime(num: int) -> bool:
    if not isinstance(num, int):
        raise TypeError("num must be integer!")
    if num < 2:
        return False
    if num == 2:
        return True
    for i in range(2, math.ceil(num ** 0.5) + 1):
        if num % i == 0:
            return False
    return True


def bw(i: int) -> tuple[int, int, int]:
    return (255 if is_prime(i) else 0,) * 3


def image(x_size: int, y_size: int = 0):
    if y_size == 0:
        y_size = x_size
    image = Image.new("RGB", (x_size, y_size), "black")
    for y in range(y_size):
        print(y * x_size)
        for x in range(x_size):
            image.putpixel((x, y), bw(x + y * x_size + 1))
    image.save("primes.png")

if __name__ == "__main__":
    image(1000)

Answer 1

There are some things I would comment on, outside of what DeathIncarnate said in their answer.

---

The is_prime is fine, I would change two small details:

1. It is uncommon in Python to check the type of the input value (if not isinstance(num, int)). If the input is not an int, then all sorts of things are going wrong in the calling function. Is it really the responsibility of this function to avoid those? You could add an assert, which can be disabled if necessary. But mostly you should rely on the static type checking, which is what you have the type annotations for. This happens outside of executing the Python code, and therefore does not influence the execution speed.

2. The range range(2, math.ceil(num ** 0.5) + 1) is too wide. If num**0.5 is not an integer, then the square of its ceil will be larger than num, and therefore does not need to be checked. You should use math.floor instead.

   But it generally contains too many numbers: if you can't divide by 2, you won't be able to divide by 4, 6, 8, ... This range could skip all even numbers. In fact, this range should skip all non-prime numbers, which is where my comment under your question hinted at.

The sieve of Erastothenes would be the standard way of producing a list of all primes up to a given number. It checks each new number to not be divisible by the primes found so far (in a manner of speaking, it actually doesn't use any divisions at all). In your case, you need a list of all primes up to and including x_size * y_size. Let's call this list primes (we'll come back to this later).

---

I agree with DeathIncarnate, you don't need an RGB image here, you just need a gray-scale image, which has a single channel. You can write a PNG of such a gray-scale image just fine. Because your code will be simpler, it is preferable.

To represent the image, I would use a NumPy array. Because NumPy arrays are easier to work with and more efficient to index than PIL.Image objects. Converting between the two is trivial, typically no data needs to be copied.

Earlier I referred primes, the list of primes. This list can be used to index into a NumPy array:

import numpy as np

# ...

img = np.zeros((y_size, x_size), dtype=np.uint8)
img.reshape(-1)[primes] = 255

That's it, that's all you need to do to create your image.

...well, in your code you start counting pixels from 1, not 0 as I did there. To subtract 1 from your primes, you can convert the list to a NumPy array first:

img.reshape(-1)[np.array(primes) - 1] = 255

Now, to save the image using Pillow, you can cast to a PIL.Image:

Image.fromarray(img).save('primes.png')

---

I wanted to try these ideas out. rdesparbes posted code with a better prime list generation similar to ideas I wrote above. This reduced the computation time on my machine from the 2.8s of OP's code to 0.63s. Using a NumPy array instead of a PIL.Image object, the time went down about 12%, to 0.55s. Implementing an actual sieve of Erastothenes (which doesn't do any divisions at all) reduced the time to 1/5th of that, 0.11s. This is about 25 times faster than OP's code.

This is the code I put together for testing. It's a bit messy, sorry for that, don't pay attention to the style, function names, or the functions themselves writing to file.

Note that about half of the 0.11s is converting the NumPy array to PIL.Image, and saving it to file.

import math
import time

import numpy as np
from PIL import Image


# --- OP ---


def is_prime(num: int) -> bool:
    if not isinstance(num, int):
        raise TypeError("num must be integer!")
    if num < 2:
        return False
    if num == 2:
        return True
    for i in range(2, math.ceil(num ** 0.5) + 1):
        if num % i == 0:
            return False
    return True


def bw(i: int) -> tuple[int, int, int]:
    return (255 if is_prime(i) else 0,) * 3


def original_image(x_size: int, y_size: int = 0):
    if y_size == 0:
        y_size = x_size
    image = Image.new("RGB", (x_size, y_size), "black")
    for y in range(y_size):
        for x in range(x_size):
            image.putpixel((x, y), bw(x + y * x_size + 1))
    image.save("original_image.png")


# --- RDESPARBES https://codereview.stackexchange.com/a/284350/151754 ---


def compute_primes(maximum: int) -> list[int]:
    # updated by Cris to not test even numbers, and a few other tweaks
    if maximum <= 2:
        return []
    primes = [2]

    def is_prime(number_to_test_: int) -> bool:
        max_threshold = int(number_to_test_ ** 0.5)
        for prime in primes:
            if prime > max_threshold:
                break
            if number_to_test_ % prime == 0:
                return False
        return True

    for number_to_test in range(3, maximum, 2):
        if is_prime(number_to_test):
            primes.append(number_to_test)
    return primes


def rdesparbes_image(x_size: int, y_size: int | None = None) -> None:
    if y_size is None:
        y_size = x_size
    primes: list[int] = compute_primes(maximum=x_size * y_size + 1)  # fixed upper limit!
    image = Image.new("RGB", (x_size, y_size), "black")
    for prime in primes:
        y, x = divmod(prime - 1, x_size)
        image.putpixel((x, y), (255, 255, 255))
    image.save("rdesparbes_image.png")


# --- CRIS #1: like above, but using NumPy array ---


def better_image(x_size: int, y_size: int | None = None) -> None:
    if y_size is None:
        y_size = x_size
    primes = compute_primes(x_size * y_size + 1)
    img = np.zeros((y_size, x_size), dtype=np.uint8)
    img.reshape(-1)[np.array(primes) - 1] = 255
    Image.fromarray(img).save('better_image.png')
    

# --- CRIS #2: sieve of Erastothenes ---


def sieve_image(x_size: int, y_size: int | None = None) -> None:
    if y_size is None:
        y_size = x_size
    primes = np.ones(x_size * y_size + 1, dtype=np.bool_)
    primes[0:2] = False
    for ii in range(2, len(primes)):
        if primes[ii]:
            primes[ii**2::ii] = False
    primes = np.reshape(primes[1:], (y_size, x_size)).astype(np.uint8)
    primes *= 255
    Image.fromarray(primes).save('sieve_image.png')


# --- MAIN ---


if __name__ == "__main__":
    t = time.time()
    original_image(1000)
    print(f"{time.time() - t:0.3f}s")

    t = time.time()
    rdesparbes_image(1000)
    print(f"{time.time() - t:0.3f}s")
    
    t = time.time()
    better_image(1000)
    print(f"{time.time() - t:0.3f}s")

    t = time.time()
    sieve_image(1000)
    print(f"{time.time() - t:0.3f}s")

Answer 2

Now, I will preface this by saying I've never used pillow before, so this is likely not the most effective means of doing this.

Docstrings

You have used type annotations well in your functions which is a good start, but you may want to look into adding docstrings to describe their function.

Arguments

Why do you default y_size to 0, as your not-entered? It is more idiomatic to use None (you may choose -1 if you want to avoid None) as a not supplied value rather than something which is potentially "valid". Consider if I was automating using this I could imagine a scenario where I requested a 0 element image (of course if I requested a -1 sized image that should probably error out)

Also, you don't allow the user to write to a file of their choosing. You should probably pass that as an argument through to the function to avoid users having to manually update their code or accidentally overwriting different sized images.

Shadowing

The main image variable inside your function has the same name as the function itself which is not a good thing and could lead to confusion. It is common practice to give functions names which are verbs such as generate_image, even this is a little generic and you might want generate_primes_image

Image Type

You create a full RGB image while pillow offers a convenient bitwise image which seems to be what you want. Additionally, with the bitmap we would no longer need bw and simply is_prime can be our generator.

The pillow documentation also notes:

Image.putpixel(xy, value)[source]

Modifies the pixel at the given position. The color is given as a single numerical value for single-band images, and a tuple for multi-band images. In addition to this, RGB and RGBA tuples are accepted for P and PA images.

Note that this method is relatively slow. For more extensive changes, use paste() or the ImageDraw module instead.

So we probably might want to look at other means of generating it.

Using generators

Looking into the documentation it seems we can generate our image more efficiently by running the loops as a generator:

def generate_primes_image(x_size: int, y_size: int = None, output_file: str = "primes.png"):
    """ Generates a pixel representation of the primes as a 2d image saved to 
        `output_file`     
    """
    if y_size is None:
        y_size = x_size

    # Using 255*is_prime because this packs to length-8 bytes
    # mode is greyscale ("L") to match 
    u = (is_prime(x)*255 for x in range(1, (x_size*y_size)+1))
    image = Image.frombytes("L", (x_size, y_size), bytes(u))
    image.save(output_file)

This should actually generate the data through C all at once which should be more efficient.

You could also look at the Sieve of Erastothenes (as suggested in the comments) or other prime generating methods and use the ImageDraw module to modify chunks of the image as you find the elements which are not prime.

Answer 3

Style and design

if not isinstance(num, int):
    raise TypeError("num must be integer!")

This piece of code is not really necessary: the type annotation makes it clear that you expect an int here. If you are worried about what the user of this function could provide as an argument, you can add a docstring, or use a static type checker like mypy

def image(x_size: int, y_size: int = 0):
    if y_size == 0:
        y_size = x_size

• This behavior could be described in a docstring, as it is not obvious from the signature of the function itself.
• Also, on a more general note, good job on adding type annotations. But it's a good practice to also indicate that a function does not return anything with -> None:.
• image is a confusing name for a function. Usually, functions start with a verb. This function's name conflicts with the name of the actual Image.Image instance. compute_image or generate_prime_image would probably be a better name for this function.

Split responsibilities

image.save("primes.png")

This part makes the image function hard to reuse in another context, and should probably be moved elsewhere.

Performance

This script is slow for two main reasons:

• The print statement, which makes expensive system calls
• Primes that have already been computed could be reused to compute new primes (see Sieve of Eratosthenes)

Refactoring

Here is another version of the script that shows the same behavior, but is apparently faster because prime numbers are saved and reused to compute new ones:

from __future__ import annotations

import math

from PIL import Image


def compute_primes(maximum: int) -> List[int]:
    """Compute the list of the prime number

    :param maximum: Upper bound of the list, not included
    :return: The prime numbers, in ascending order
    """
    if maximum <= 2:
        return []
    primes: List[int] = []

    def is_prime(number_to_test_: int) -> bool:
        max_threshold: int = math.ceil(number_to_test_ ** 0.5) + 1
        for prime in primes:
            if prime >= max_threshold:
                break
            if number_to_test_ % prime == 0:
                return False
        return True

    for number_to_test in range(2, maximum):
        if is_prime(number_to_test):
            primes.append(number_to_test)
    return primes


def generate_prime_image(x_size: int, y_size: int | None = None) -> Image.Image:
    if y_size is None:
        y_size = x_size
    primes: list[int] = compute_primes(maximum=x_size * y_size)
    image = Image.new("RGB", (x_size, y_size), "black")
    for prime in primes:
        y, x = divmod(prime - 1, x_size)
        image.putpixel((x, y), (255, 255, 255))
    return image


def main() -> None:
    image = generate_prime_image(1000)
    image.save("primes.png")


if __name__ == "__main__":
    main()

Answer 4

I'm not a Python "expert", but, I hope you appreciate what I have to say.

---

Besides everything that was said, I do see 2 issues with your code:

1 - You have a magic number

You have the following lines:

if __name__ == "__main__":
    image(1000)

You could use the argparse module to receive the size as an argument.

Here's my example on how it can be implemented:

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        prog="Some Program name here",
        description="Generates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't."
    )
    
    parser.add_argument("-s", "--size", type=int, default=1000)
    args = parser.parse_args()
    
    image(args.size)

---

2 - The image name and path are hardcoded

This makes it pretty inflexible to use.

You can change your image function to have the following:

def image(x_size: int, y_size: int = 0, fp = "primes.png"):
    [...]
    image.save(fp)

Then, you can modify the code I gave you before, to allow passing a file name:

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        prog="Some Program name here",
        description="Generates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't."
    )
    
    parser.add_argument("-s", "--size", type=int, default=1000)
    parser.add_argument("-f", "--file", type=str, default="primes.png")
    args = parser.parse_args()
    
    image(args.size, fp=args.filename)

---

This basically allows you to actually make something useful out of that if __name__ == "__main__":.

You can just run it like this: python3 file.py --size=2500 --file="huge.jpg".