Creates test images with a grid of simple, Zener-card-like icons

zener.py

import sys
import csv
import random
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import numpy as np

# Constants
COLORS = ['red', 'blue', 'green', 'orange', 'purple']
SHAPES = ['circle', 'square', 'triangle', 'plus', 'diamond']

def draw_random_icon(draw, x, y, cell_size):
    icon_size = cell_size // 3
    left = x + (cell_size - icon_size) // 2
    top = y + (cell_size - icon_size) // 2
    right = left + icon_size
    bottom = top + icon_size
    center_x = (left + right)/2
    center_y = (top + bottom)/2
    shape = random.choice(SHAPES)
    color = random.choice(COLORS)

    if shape == 'circle':
        draw.ellipse([left, top, right, bottom], fill=color)
    elif shape == 'square':
        draw.rectangle([left, top, right, bottom], fill=color)
    elif shape == 'plus':
        # two long overlapping rectangles
        draw.rectangle([
            center_x-5, top,
            center_x+5, bottom
        ], fill=color)
        draw.rectangle([
            left, center_y-5,
            right, center_y+5
        ], fill=color)
    elif shape == 'triangle':
        draw.polygon([ 
            (center_x, top),
            (left, bottom),
            (right, bottom)
        ], fill=color)
    elif shape == 'diamond':
        draw.polygon([
            (center_x, top),
            (left, center_y),
            (center_x, bottom),
            (right, center_y),            
        ], fill=color)
    return f"{color} {shape}"

# Draw the grid and icons
def random_grid(grid_size, image_size=512):
    image = Image.new('RGB', (image_size, image_size), 'white')
    draw = ImageDraw.Draw(image)
    cell_size = image_size // grid_size
    grid_data = []
    for row in range(grid_size):
        row_data = []
        for col in range(grid_size):
            x = col * cell_size
            y = row * cell_size
            description = draw_random_icon(draw, x, y, cell_size)
            row_data.append(description)
        grid_data.append(row_data)
    return image, grid_data



if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("Usage: python script.py <grid_size> <output_filename>")
        sys.exit(1)
    
    grid_size = int(sys.argv[1])
    output_filename = sys.argv[2]

    image, data = random_grid(grid_size)
    image.save(output_filename)

    csv_filename = output_filename.rsplit('.', 1)[0] + ".csv"
    with open(csv_filename, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(data)

    print(f"Grid image saved as {output_filename} and data saved as {csv_filename}")