PythonのPILで２次元配列データから画像を作成するプログラム

下記のようなデータから
[[0, 1, 0, 1],
[1, 1, 0, 0],
[1, 1, 0, 1],
[0, 1, 1, 0]]

下記のような画像を作成するプログラムです。

from PIL import Image, ImageDraw

def draw_square(row, column, color):
    x1 = column * square_size
    y1 = row * square_size
    x2 = column * square_size + square_size
    y2 = row * square_size + square_size
    draw.rectangle((x1, y1, x2, y2), fill=color)

data = [[0, 1, 0, 1],
        [1, 1, 0, 0],
        [1, 1, 0, 1],
        [0, 1, 1, 0]]
data_size_width = len(data[0])
data_size_height = len(data)

square_size = 20

image_size_width = square_size * data_size_width
image_size_height = square_size * data_size_height

im = Image.new("RGB", (image_size_width, image_size_height), 'white')
draw = ImageDraw.Draw(im)
for row in range(data_size_height):
    for column in range(data_size_width):
        if data[row][column] == 0:
            draw_square(row, column, 'white')
        else:
            draw_square(row, column, 'black')
im.save('a.jpg')

今ライフゲームのプログラムを作成していて、結果を画像で保存したいので、そのために作成したものです。