Licence

GraphicsWithPythonAndNumpy_1_0_0.ipynb - Graphics with Python and Numpy notebook

Written in 2019 by Andrew Baldwin (twitter:baldand)

To the extent possible under law, the author(s) have dedicated all copyright and related and neighboring rights to this software to the public domain worldwide. This software is distributed without any warranty. See for CC0 Public Domain Dedication.

Graphics with Python and Numpy

This notebook shows how to make simple animated graphics & videos using Python and Numpy and PIL.

The intent is that all the techniques and code show can be usefully used in any kind of environment, from a server, to a Raspberry Pi. A GPU is not required.

Another target is that there is no hidden magic libraries. Apart from the standard libraries (numpy and PIL) - all the functions needed for working with and generating are defined here in this notebook.

Also, we will not be generating images with high level functions for drawing shapes on a canvas. Instead everything will be using lower-level mathematical approaches (similar to those you might use when wrinting a shader in OpenGL or WebGL).

If you are viewing this from Jupyter, you can run the cells in order to see the results generated by your python environment.

To get a better understanding of how the code works, try changing some of the numbers to see how they influence the output images.

If you are looking at a static version of the notebook, you can see what would result from running the code.

Libraries

First we need to import the python libraries we will use later

In [1]:
# Standard python libraries
import math # For math functions
import time # For time related things
from pathlib import Path # For file things
 
# Add on libraries that may need to be installed
import numpy as np # For working with our image data
from PIL import Image # For creating viewable images from our data 
from IPython import display # For showing the images

Setup

We start by creating some functions to help us work with images.

In [2]:
def nptopil(image_array):
    """ 
    Convert from an image in a numpy array intto a PIL Image
    Array should have shape (height,width,channels), dtype float32, range 0-1
    Output is PIL Image, unsigned 8bit integer, range 0-255 
    """
    return Image.fromarray(np.squeeze(np.clip(image_array, 0.0, 1.0)*255).astype('B')) 

def show(npimg, handle=None):
    """
    Display an np img in notebook
    Returns a handle which can be used to update an already displayed image
    """
    pilimg = nptopil(npimg)
    if handle is not None:
        handle.update(pilimg) # Used for animation later
    else:
        handle = display.display(pilimg, display_id=True)
    return handle

Basics

Now we are ready to learn some basic principles for working with images.

All our images will be created as 3-dimensional numpy arrays of floating point numbers.

These arrays will have the shape (height in pixels, width in pixels, number of channels).

For greyscale images, the channel value is 1 - a single value gives the level.

For colour images, the channel count is 3 - for red, green and blue levels of the pixel.

Every location in the array will store the brightness of the channel for a pixel, in the range 0.0-1.0

Let's create an array of black pixels:

In [3]:
height, width = 100, 200 # This will be the size in pixels 
all_black = np.zeros(shape=(height, width, 1), dtype=np.float32) # Create a greyscale image of all black pixels
show(all_black); # Now show it as the output of this cell. (";" at the end suppresses an extra print)

We can also create an array of white pixels:

In [4]:
all_white = np.ones(shape=(height, width, 1), dtype=np.float32)
show(all_white);

We can create new images from existing ones, for example by multiplying the values of all the pixels by a constant:

In [5]:
all_grey = all_white * 0.5
show(all_grey);

We can directly change the values of pixels in an array, either one at a time, or in a block:

In [6]:
# Directly change the values of pixels like this: 
# im[from_top,from_left] = new_value for pixel
# im[top:bottom,left:right] = new_value for block

all_grey[10,10] = 1.0 # One white pixel
all_grey[25:75, 25:75]   = 0.25 # Dark grey square on the left

# We can also multiply and add to regions in a similar way
all_grey[25:75, 125:175] *= 2.0 # Make square on white white again
all_grey[40:60, 50:150] -= 0.25 # Make darker

show(all_grey);
In [7]:
# We can also save any image to a file like this:
nptopil(all_grey).save("image1.png")
# In JupyterLab you should be able to see the file in the left sidebar and open it
In [8]:
# Colour images are similar
# For each pixel there are three values for the brightness of red, green and blue
col_im = np.zeros(shape=(height, width, 3), dtype=np.float32)
col_im[25:75, 50:150] = (1.0,0.5,0.25) # An orange rectangle (Red=1, Green=0.5, Blue=0.25)
show(col_im);
In [9]:
# We can cut out parts of images
left = col_im[:,:100] # Take the left side
show(left);
In [10]:
right = col_im[:,100:]
show(right);
In [11]:
# We can combine images
mixed = left + right*0.5 # Half the brightness
show(mixed);

Making images with maths

Now you've seen what we are working with. The next question is how will we generate images.

Rather than changing specific pixels, which would fix the size of our images, we want to work in a way that lets us render images at different sizes.

To achieve that, we are going to define our images as functions. These will take arguments which can be in the form of images (arrays) giving an input value for each output pixel. The function will then calculate the output values from the inputs.

The primary input values we will work with will be a pair of coordinates:

  • x, for the horizontal position, increasing from 0 on the left to 1 on the right
  • y, for the vertical position, increasing from 0 on the bottom to 1 on the top

Once we have our image generating functions, we will use a rendering function to call those with suitable input values to render our output images.

Now we need to create some more helper functions for this.

In [12]:
def coords(height, width):
    """
    Return y and x coordinate arrays both with shape (height, width)
    """
    x = np.ones(shape=(height, width, 1), dtype=np.float32)
    x *= np.arange(width, dtype=np.float32).reshape(1, width, 1) / width + 0.5/width
    y = np.ones(shape=(height, width, 1), dtype=np.float32)
    y *= np.arange(height-1, -1, -1, dtype=np.float32).reshape(height, 1, 1) / height + 0.5/height
    return y, x

def render(height, width, image_fn, *params, t=None, ms=(4,4), msthresh=0.1):
    """
    Render image_fn at the given size and return as array
    Multisample edges with a grid of subpixels defined by ms (skipped if ms=(1,1))
    """
    y, x = coords(height, width)
    msh, msw = ms
    msamp = msh * msw
    if t is None:
        image_array = image_fn(y, x, *params)
    else:
        image_array = image_fn(y, x, t, *params)
    # Multi-sample edge pixels
    if msamp > 1:
        # Find edge pixels
        grads = np.empty(shape=(image_array.shape[0], image_array.shape[1]))
        diff = np.sum(np.abs(image_array[:-1,:-1]-image_array[:-1:,1:]),axis=2)
        diff += np.sum(np.abs(image_array[:-1,:-1]-image_array[1:,:-1]),axis=2)
        grads[:-1,:-1] = diff
        grads[:-1,1:] += diff
        grads[1:,:-1] += diff
        grads[1:,1:] += diff
        edges = np.argwhere(grads > msthresh)
        # Generate grid of sub-coords for edge pixels
        yms = np.empty(shape=(edges.shape[0],msamp), dtype=np.float32)
        oy,ox = np.mgrid[:msh,:msw]
        yms[:,:] = y[edges[:,0],edges[:,1]] + ((2.*oy.reshape(1,msamp)-(msh-1))*(1.0/(msh*height)))
        xms = np.empty(shape=(edges.shape[0],msamp), dtype=np.float32)
        xms[:,:] = x[edges[:,0],edges[:,1]] + ((2.*ox.reshape(1,msamp)-(msw-1))*(1.0/(msw*width)))
        # Sample each edge pixel multiple times
        if t is None:
            ms = image_fn(yms.reshape(edges.shape[0],msamp,1), xms.reshape(edges.shape[0],msamp,1), *params)
        else:
            ms = image_fn(yms.reshape(edges.shape[0],msamp,1), xms.reshape(edges.shape[0],msamp,1), t, *params)
        image_array[edges[:,0],edges[:,1]] = np.sum(ms, axis=1)*(1.0/msamp)
    return image_array

def vec(*vals):
    # This ensures colour values are float32 arrays
    return np.array(vals, dtype=np.float32)

Gradients

Let's start with a very simple image function:

In [13]:
def grad(y, x):
    # Draw a horizontal gradient
    # Just return the y coordinates as pixel values
    return y

show(render(200, 200, grad));

This returns the coordinates as a brightness to create a smooth gradient from black to white.

Next let's introduce a function called mix which we can use to extend this to coloured gradients (that might be useful for backgrounds later):

In [14]:
def mix(x, y, a):
    af = np.clip(np.float32(a),0.0,1.0)
    return x * (1-a) + y * a

def hgrad(y, x, col_b=vec(0,0,0), col_t=vec(1,1,1)):
    """
    A horizontal gradient 
    """
    return mix(col_b, col_t, y)

def vgrad(y, x, col_b=vec(0,0,0), col_t=vec(1,1,1)):
    """
    A vertical gradient 
    """
    return mix(col_b, col_t, x)

cyan = vec(0,1,1) # (red, green, blue) levels between 0-1
blue = vec(0,0,1)
show(render(200, 200, hgrad, cyan, blue));

Step, Squares and rotation

Now we introduce another important function called step.

It takes an edge parameter, and a value. If the value is greater than or equal to the edge, it returns 1. Otherwise it returns 0.

In [15]:
def step(edge, v):
    """
    Returns 1 if v is greater than or equal to edge. Otherwise returns 0
    """
    return (v>=edge)*np.float32(1.0)

def wall_scene(y, x, edge):
    cyan = vec(0,1,1)
    blue = vec(0,0,1)
    bg = hgrad(y, x, cyan, blue) # Reuse as background
    black = vec(0,0,0)
    grey = vec(0.5,0.5,0.5)
    fg = vgrad(y, x, grey, black) # Also use vgrad as foreground
    
    w = step(edge, x) # The step function, makes a hard edge
    
    return mix(bg, fg, w) # Use w as mask to select foreground or background

show(render(200, 200, wall_scene, 0.5)); # Try changing 0.5 to 0.75 or 0.25

To invert the step function, just give the arguments in the opposite order

In [16]:
def wall_inverted(y, x, edge):
    cyan = vec(0,1,1)
    blue = vec(0,0,1)
    bg = hgrad(y, x, cyan, blue) # Reuse as background
    black = vec(0,0,0)
    grey = vec(0.5,0.5,0.5)
    fg = vgrad(y, x, black, grey) # Also use vgrad as foreground
    
    w = step(x, edge) # The step function (inverted), makes a hard edge
    
    return mix(bg, fg, w) # Use w as mask to select foreground or background

show(render(200, 200, wall_inverted, 0.5)); # Try changing 0.5 to 0.75 or 0.25

Because step returns only 0 or 1, and 0 0, 1 0 and 0 1 = 0, but 1 1 = 1, we can combine step outputs with multiplication.

In [17]:
def square(y, x, size=0.5):
    """
    Draw a square centred on (0,0)
    """
    hsize = size*0.5
    return step(-hsize, x) * step(x, hsize) * step(-hsize, y) * step(y, hsize) 

def square_scene(y, x):
    cyan = vec(0,1,1)
    blue = vec(0,0,1)
    bg = hgrad(y, x, cyan, blue) # Reuse as background
    white = vec(1,1,1)
    grey = vec(0.5,0.5,0.5)
    fg = hgrad(y, x, grey, white) # Also use vgrad as foreground
    
    sq = square(y-0.5, x-0.5, 0.5) # Try changing these numbers
    
    return mix(bg, fg, sq) # Use sq as mask to select foreground or background

show(render(400, 400, square_scene, ms=(1,1)));

Squares are nice, and perhaps you can already see how to change the size and position of this one.

But what about if you would like to rotate it?

Our square is drawn using the y and x coordinate arrays. A rotated version of the square would still have y and x arrays, but the values would be different.

We can transform our original y and x to the rotated y and x with a simple function. This function can be used to rotate the output of any of the functions we have seen so far.

In [18]:
def rotate(y, x, angle):
    """
    Return y & x coordinates rotated by angle radians
    """
    sx = np.sin(angle)
    cx = np.cos(angle)
    return x * sx + y * cx, x * cx - y * sx

def square_rotated_scene(y, x, angle):
    cyan = vec(0,1,1)
    blue = vec(0,0,1)
    bg = hgrad(y, x, cyan, blue)
    white = vec(1,1,1)
    grey = vec(0.5,0.5,0.5)
    fg = hgrad(y, x, grey, white) # Also use vgrad as foreground
    
    yr, xr = rotate(y-0.5, x-0.5, angle)
    sq = square(yr, xr, 0.5) # Try changing these numbers
    
    return mix(bg, fg, sq) # Use sq as mask to select foreground or background

# Try changing the angle here. Remember 2*math.pi = 360 degrees
# Question: does it rotate also the gradient filling the square?
i = render(400, 400, square_rotated_scene, math.pi/3.)
nptopil(i).save("square.png")
show(i);

Length and Angle

Now lets introduce the length function, which uses Pythagoras' theorem about right angled triangles to calculate the length of a 2d vector.

Using length, we can make circles.

In [19]:
def length(y, x):
    """
    Return the length (magnitude) of the vector (y,x)
    """
    return np.sqrt(x*x + y*y)

def circle(y, x, radius=0.25):
    """
    Return a circular mask centred at (0,0)
    """
    return step(length(y, x), radius)
    
def rgrad(y, x, col_c=vec(0,0,0), col_e=vec(1,1,1)):
    """
    Radial gradient
    """
    return mix(col_c, col_e, length(y, x))

def circle_scene(y, x):
    red = vec(0.75,0,0)
    magenta = vec(0.75,0,0.75)
    bg = hgrad(y, x, red, magenta)
    orange = vec(1.0,0.5,0.0)
    yellow = vec(1,1,0.5)
    sunx = 0.5
    suny = 0.4
    sunr = 0.3
    fg = rgrad((y - suny)/sunr, (x - sunx)/sunr, yellow, orange) # Sun
    
    cir = circle(y - suny, x - sunx, sunr) 
    
    return mix(bg, fg, cir) # Use cir as mask to select foreground or background

show(render(400, 400, circle_scene));

The length function calculates the length of the vector (y,x).

We can also calculate the angle of the same vector using trigonometry, namely the arctan function.

Let's define a function angle to do that, and then make use of it.

In [20]:
def angle(y, x):
    # Will return angle of the vector (y,x) in radians between 0 and 2*pi (360 degrees)
    return np.arctan2(x, y)+math.pi

def pac(y, x):
    ang = angle(y-0.5, x-0.5) # We will use the angle to cut out part of a circle
    deg = 180*ang/math.pi
    mouth = 1 - step(240, deg)*step(deg, 300)
    
    eye = 1 - circle(x-0.5, y-0.7, 0.05)
    body = circle(x-0.5, y-0.5, 0.35)
    yellow = vec(1.0, 1.0, 0.0)
    bodycol = rgrad((x-0.5)/0.35, (y-0.5)/0.35, yellow, yellow*0.7)
    food = circle(x-0.9, y-0.5, 0.05)
    blue = vec(0,0,0.9)
    walls = (step(y, 0.05) + step(0.95, y)) * blue
    return eye * mouth * body * bodycol + food + walls
    
show(render(400, 400, pac));

Different shapes

Using angle and length together, we can create a more general function that can make very many interesting shapes.

In [21]:
def starflower(y, x, radius=0.25, sides=4.0, rotate=0.0, bend=1.0):
    """
    Draw a shape centred at (0,0)
    """
    a = angle(y, x) + rotate
    l = length(y, x)
    b = 2*math.pi / sides
    m = b*np.floor(0.5 + a/b) - a
    c = np.cos(m)**bend
    return step(l * c, radius)
 
def box(y, x):
    dark_blue = vec(0,0,0.1)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)

    # With no arguments, it is another way to draw squares
    shape = starflower(y-0.5, x-0.5, rotate=math.pi/3.)
    red = (1,0,0)
    return mix(bg, red, shape)

show(render(400, 400, box));
In [22]:
def tri(y, x):
    dark_blue = vec(0,0,0.1)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)

    # But, you can change the number of sides. Here is a triangle
    shape = starflower(y-0.5, x-0.5, radius=0.2, sides=3, rotate=math.pi/5.)
    olive = (0.4, 0.5, 0.25)

    return mix(bg, olive, shape)

show(render(400, 400, tri));
In [23]:
def pent(y, x):
    dark_blue = vec(0,0,0.1)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)

    # Pentagon
    shape = starflower(y-0.5, x-0.5, radius=0.3, sides=5, rotate=2.*math.pi/3.)
    turquoise = (0.2, 0.65, 0.6)

    return mix(bg, turquoise, shape)

show(render(400, 400, pent));
In [24]:
def star(y, x):
    dark_blue = vec(0,0,0.1)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)

    # Pentagon
    shape = starflower(y-0.5, x-0.5, radius=0.2, sides=5, bend=4.0, rotate=4.*math.pi/3.)
    staryellow = (1.0, 1.0, 0.4)

    return mix(bg, staryellow, shape)

show(render(400, 400, star));
In [25]:
def flower(y, x):
    dark_blue = vec(0,0,0.1)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)

    # Pentagon
    shape = starflower(y-0.5, x-0.5, radius=0.4, sides=5, bend=-6.0, rotate=5.*math.pi/3.)
    pink = (1.0, 0.2, 0.6)

    return mix(bg, pink, shape)

show(render(400, 400, flower));

Repetition

In a similar way that we rotate y and x earlier, we can also transform them into a grid in order to create more objects.

In [26]:
def grid(h, w, y, x):
    """
    Subdivide y and x into a grid of h x w cells
    Return new y and x, and also yi and xi - the integer indexes of the cells  
    """
    xf = (x*w) % 1
    yf = (y*h) % 1
    xi = np.floor(x*w)
    yi = np.floor(y*h)
    return yf, xf, yi, xi

def shapegrid(y, x):
    dark_blue = vec(0,0,0.01)
    blue = vec(0,0,0.5)
    bg = hgrad(y, x, dark_blue, blue)
    
    g3_y, g3_x, g3_yi, g3_xi = grid(3, 3, y, x)
    
    shape = starflower(g3_y-0.5, g3_x-0.5, sides=g3_yi+5.0, radius=0.15, bend=(g3_xi-1.0)*4.0+1.0)
    red = (1.0, 0.0, 0.0)
    green = (0.0, 1.0, 0.0)
    mid_blue = (0.0, 0.0, 0.5)
    colour = mix(green, red, g3_xi/2.0) + mix(mid_blue, red, g3_yi/2.0)
    
    return mix(bg, colour, shape)

show(render(400, 400, shapegrid));

Animation

Animating requires time - a time value, which can be used to modify things that should change.

Then, instead of creating just one still image, many images need to be drawn, each time using a different value for time.

This can be done with a loop to calculate the time value, then draw and show the new frame

In [27]:
def spinning(y, x, t): # t is time in seconds
    # This is the function to generate a single frame
    black = vec(0,0,0)
    blue = vec(0,0,0.5)
    bg = rgrad((y-0.5)/0.5, (x-0.5)/0.5, blue, black)

    g3_y, g3_x, g3_yi, g3_xi  = grid(5, 5, y, x)
    phase = math.pi*t*(10+g3_xi+g3_yi)/21 # Each shape has different frequency, align every 21 seconds
    shapetype = 4*np.sin(phase) # Change shape
    shape = starflower(g3_y-0.5, g3_x-0.5, radius=0.2, sides=5, rotate=phase, bend=shapetype)
    shade = np.sin(phase)*vec(0.0, 0.4, -0.3) + vec(1.0, 0.6, 0.3) # Change colour
    return mix(bg, shade, shape)

# This is the function to animate
def animate(height, width, drawfn, duration, step):
    handle = None
    t = 0
    while t < duration:
        img = render(height, width, drawfn, t=t, ms=(2,2))
        time.sleep(step)
        handle = show(img, handle)
        t += step
        
animate(200, 200, spinning, 21, 1/30) # Use the stop button above to interrupt the kernel

Rendering videos

Once you have a time dependant image function, it can be used to render higher resolution image files for each time which can then be used to encode a video file.

Encoding can be done by calling ffmpeg as an external process.

Finally the resulting video can be loaded back into the notebook to preview.

In [ ]:
def render_video(framedir, height, width, drawfn, duration, step):
    Path(framedir).mkdir(exist_ok=True)
    y, x = coords(height, width)
    t = 0
    framenum = 0
    h = display.display("Rendering %.02f/%.02f"%(t,duration), display_id=True)
    while t < duration:
        #display.clear_output()
        h.update("Rendering %.02f/%.02f"%(t,duration))
        frame = render(height, width, drawfn, t=t, ms=(8,8)) # Use more multisampling for smoother edges
        nptopil(frame).save(framedir+"/frame_%04d.png"%framenum)
        t += step
        framenum += 1
    h.update("Rendering complete")
        
render_video("spinning",720,720,spinning,21,1/30) # Good size for twitter
In [ ]:
!ffmpeg -y -r 30 -i spinning/frame_%04d.png -an -pix_fmt yuv420p -r 30 -b:v 2M spinning/spinning.mp4
In [28]:
display.Video("spinning/spinning.mp4")
Out[28]:
In [ ]: