Testing halftone speeds

a_test.py

# Comparing answers to this Stack Overflow question:
# https://stackoverflow.com/q/58791296/

from original import process_tones, halftone
from tstanisl import process_tones2, halftone_fast
from jirassimok import tones_as_array, np_halftone

def full_halftone(img):
    return halftone(img, process_tones())

def full_halftone_fast(img):
    return halftone_fast(img, *process_tones2())

def full_np_halftone(img):
    return np_halftone(img, tones_as_array())

def randimage(h, w):
    from numpy import random, uint8
    return random.randint(0, 256, (h, w), uint8)

if __name__ == '__main__':
    import timeit
    n = 100
    setup = ("from __main__ import randimage, full_halftone, full_halftone_fast, full_np_halftone\n"
             f"img = randimage({n}, {n})")

    # Run the original halftone function 10 times
    orig = timeit.repeat(f'full_halftone(img)',
                         setup=setup, repeat=10, number=1)

    # Run the other two in r sets of n runs each
    r, n = 10, 1000 # trials per image
    fast  = timeit.repeat(f'full_halftone_fast(img)',
                          setup=setup, repeat=r, number=n)
    fast = (t / n for t in fast)

    npht  = timeit.repeat(f'full_np_halftone(img)',
                          setup=setup, repeat=r, number=n)
    npht = (t / n for t in npht)

    print("Best times")
    print(f"halftone:      {min(orig):.15f}")
    print(f"np_halftone:   {min(npht):.15f}")
    print(f"halftone_fast: {min(fast):.15f}")

jirassimok.py

"""
Copyright 2019 jirassimok on Stack Overflow (https://stackoverflow.com/users/7389264/jirassimok)

Source: https://stackoverflow.com/a/58792592/

CC-BY SA 4.0 (https://creativecommons.org/licenses/by-sa/4.0/)
"""
import numpy as np
from original import TONES

def tones_as_array():
   global TONES
   num_tones = len(TONES)
   if num_tones <= 0:
      raise ValueError("no tones")
   xy = len(TONES[0])
   x = y = int(np.sqrt(xy))
   return 255 * np.array(TONES).reshape((num_tones, x, y))

def np_halftone(grayscale, tones):
    """Convert grayscale image to halftones.

    The grayscale image should have shape (h, w) and values on [0, 255].

    The tone array should have shape (n, x, y), where n is the number of tones
    available, and (x, y) is the shape of each dot matrix.

    The resulting halftone has shape (h*x, w*y).
    """
    h, w = grayscale.shape
    num_tones, x, y = tones.shape
    bins = np.linspace(0, 256, num_tones + 1)
    levels = np.digitize(grayscale, bins) - 1
    return tones[levels].swapaxes(1, 2).reshape(h * x, w * y)

original.py

"""
Copyright 2019 miguelmorin on Stack Overflow (https://stackoverflow.com/users/9251158/miguelmorin)

Source: https://stackoverflow.com/q/58791296

CC-BY SA 4.0 (https://creativecommons.org/licenses/by-sa/4.0/)
"""
import math
import time

import numpy as np

TONES = [[0, 0,
          0, 0],
         [0, 1,
          0, 0],
         [1, 1,
          0, 0],
         [1, 1,
          0, 1],
         [1, 1,
          1, 1]]

def process_tones():
    """Converts the tones above to the right shape."""
    tones_dict = dict()

    for t in TONES:
        brightness = sum(t)
        bitmap_tone = np.reshape(t, (2, 2)) * 255
        tones_dict[brightness] = bitmap_tone
    return(tones_dict)

def halftone(gray, tones_dict):
    """Generate a new image where each pixel is replaced by one with the values in tones_dict.
    """

    num_rows = gray.shape[0]
    num_cols = gray.shape[1]
    num_tones = len(tones_dict)
    tone_width = int(math.sqrt(num_tones - 1))

    output = np.zeros((num_rows * tone_width, num_cols * tone_width),
                         dtype = np.uint8)

    # Go through each pixel
    for i in range(num_rows):
        i_output = range(i * tone_width, (i + 1)* tone_width)

        for j in range(num_cols):
            j_output = range(j * tone_width, (j + 1)* tone_width)

            pixel = gray[i, j]
            brightness = int(round((num_tones - 1) * pixel / 255))

            output[np.ix_(i_output, j_output)] = tones_dict[brightness]

    return output

tstanisl.py

"""
Copyright 2019 tstanisl on Stack Overflow (https://stackoverflow.com/users/4989451/tstanisl)

Source: https://stackoverflow.com/a/58792441

CC-BY SA 4.0 (https://creativecommons.org/licenses/by-sa/4.0/)
"""
import numpy as np
from original import TONES

def process_tones2():
    tones = np.array(TONES, dtype='u1')
    size = int(np.sqrt(tones.shape[-1]))
    tones = 255 * tones.reshape(-1, size, size)
    bins = tones.sum(axis=(-2,-1), dtype=int) // size ** 2
    iperm = np.argsort(bins)
    return bins[iperm], tones[iperm]

def halftone_fast(gray, bins, tones):
    height, width = gray.shape
    tone_height, tone_width = tones.shape[-2:]
    brightness = np.round(gray / 255 * (len(tones) - 1)).astype('u1')
    binary4d = tones[brightness]
    binary4d = binary4d.transpose((0,2,1,3))
    binary = binary4d.reshape(height * tone_height, width * tone_width)
    return binary