Data set generation with Python and Matplotlib

Question

I want to use Python and Matplotlib to create a data set of images. I want to do this by transforming the plot to a numpy matrix. At the moment it takes about 0.1 second to generate 100 images of size 50x50 pixels:

Question: How can I speed up things a lot?

This is my code:

import matplotlib.pyplot as plt
import numpy as np
import time
import cv2

fig = plt.figure(frameon=False)
ax = fig.add_axes([0., 0., 1., 1.])
fig.set_size_inches((0.5,0.5))
fig_size = fig.canvas.get_width_height() + (3,)
points, = ax.plot([],[])
ax.set_xlim(0.,1.)
ax.set_ylim(0.,1.)
ax.set_axis_off()
ax.set_frame_on(False)
ax.grid(False)

def plot2mat():
    data = np.random.rand(20,2)
    points.set_xdata(data[:,0])
    points.set_ydata(data[:,1])
    fig.canvas.draw()
    fig.canvas.flush_events()
    M = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
    M = M.reshape(fig_size)
    M = cv2.cvtColor(M, cv2.COLOR_BGR2GRAY)
    return M

I = []
t0 = time.time()
runs = 1000
for k in range(runs):
    I.append(plot2mat())
print(time.time()-t0)

plt.close(fig)

for k in range(100):
    plt.subplot(10,10,k+1)
    plt.imshow(I[k])
    plt.axis("off")
plt.savefig("plot.png", dpi=400)

Answer 1

First: what is actually being plotted? Each image data M is a 50*50 array of uint8, with only two values: 127 and 255 (matplotlib makes the yellow background due to the colour map). Each line is three pixels wide with no antialiasing. There are 20 vertices making 19 lines per image.

I think that the code is way too tightly-coupled with matplotlib. The following code does not help; it's about five times slower and much more complicated, but it demonstrates that generation of your dataset is possible without matplotlib. It uses linear algebra in up to five dimensions.

import time
import typing

import numpy as np


def generate_data(
    rand: np.random.Generator,
    n_images: int = 100,
    n_lines: int = 19,
    line_width: float = 3.,
    image_width: int = 50,
    norm_method: typing.Literal['einsum_param', 'einsum_diff', 'linalg'] = 'einsum_diff',
) -> np.ndarray:
    # coordinate space in each image
    ii = np.arange(image_width)  # *50
    xy = np.stack(np.meshgrid(ii, ii), axis=-1)  # 50*50*2

    # random start/end vertices for all line segments
    vertices = rand.integers(low=0, high=image_width, size=(1 + n_lines, n_images, 2))  # 20*2
    p0 = vertices[:-1]  # 19*100*2
    p1 = vertices[1:]   # 19*100*2

    # linear parameters for singularity-free homogeneous form [abc].[xy1] = 0
    dp = p1 - p0  # 19*100*2
    c = p0[...,1]*dp[...,0] - p0[...,0]*dp[...,1]  # *19
    ab = dp[..., ::-1]*(1, -1)  # 19*2
    a = ab[..., 0]  # 19*100
    b = ab[..., 1]  # 19*100

    # matrix for projection onto the line
    denom = 1/np.einsum('ijk,ijk->ij', ab, ab)  # 19*100
    ba22 = np.broadcast_to(ab[..., ::-1, np.newaxis], ab.shape + (2,))  # 19*100*2*2
    projection = (  # 19*100*2*2
        denom[..., np.newaxis, np.newaxis] * ba22 * ba22.transpose((0, 1, 3, 2))
    ) * (
        (1, -1),
        (-1, 1),
    )

    # projected points
    xyp = (  # 19*100*50*50*2
        np.einsum('ijk,abkl->abijl', xy, projection)
    ) - (
        (c*denom)[..., np.newaxis]*ab
    )[:, :, np.newaxis, np.newaxis, :]

    # distances to line, three methods
    if norm_method == 'einsum_param':
        dist = np.abs(  # 19*100*50*50
            np.einsum('ijk,abk->abij', xy, ab)
            + c[..., np.newaxis, np.newaxis]
        )/np.hypot(a, b)[..., np.newaxis, np.newaxis]
        bound = 0.5*line_width
    elif norm_method == 'einsum_diff':
        diff = xy - xyp  # 19*100*50*50*2
        dist = np.einsum('ijklm,ijklm->ijkl', diff, diff)  # 19*100*50*50
        bound = (0.5*line_width)**2
    elif norm_method == 'linalg':
        dist = np.linalg.norm(xy - xyp, axis=-1)  # 19*100*50*50
        bound = 0.5*line_width

    # endpoint test: if 0 <= lhs <= rhs, then the projected point is on the segment
    lhs = (xyp - p0[:,:, np.newaxis, np.newaxis, :]).sum(axis=-1)  # 19*50*50
    rhs = dp.sum(axis=-1)  # *19
    r_sign = np.sign(rhs)  # *19
    lhs *= r_sign[..., np.newaxis, np.newaxis]
    rhs *= r_sign

    data = np.full(shape=(n_images, image_width, image_width), fill_value=255, dtype=np.uint8)  # 50*50
    data[
        (
            # width criteria
            (dist <= bound)
            # endpoint criteria
            & (lhs >= 0) & (lhs <= rhs[..., np.newaxis, np.newaxis])
        ).any(axis=0)
    ] = 128

    return data


def demo() -> None:
    rand = np.random.default_rng(seed=0)
    for i in range(3):
        for method in (
            'linalg', 'einsum_param', 'einsum_diff',
        ):
            start = time.perf_counter()
            generate_data(rand=rand, norm_method=method)
            dur = time.perf_counter() - start
            print(method, dur)


if __name__ == '__main__':
    demo()

Guessing at a method that would be more efficient (short of throwing everything out and writing in C), you could try ImageDraw.line (for a single line segment), or perhaps LineCollection.