CS4HS box image generator

cs4hs-boxes.py

import csv
from PIL import Image, ImageDraw

RAW_DATA = "data.csv"
GRID_SIZE = 6
BOX_SIZE = 10
IMAGE_SIZE = (GRID_SIZE * BOX_SIZE, GRID_SIZE * BOX_SIZE)


with open(RAW_DATA) as csvfile:
    reader = csv.reader(csvfile)
    next(reader) # Skip header
    submissions = []
    for data_row_number, data_row in enumerate(reader):
        # Blank grid
        rows = []
        for column in range(0, GRID_SIZE):
            rows.append([0] * GRID_SIZE)

        # Fill grid
        for row_num, row in enumerate(data_row):
            if row:
                column_values = row.split(",")
                for column_number in column_values:
                    column_number = int(column_number[-1])
                    rows[row_num][column_number - 1] = 1

        # Add grid
        submissions.append(rows)

# Create images
for submission_number, submission in enumerate(submissions):
    filename = "image-{}.png".format(submission_number)
    image = Image.new("1", IMAGE_SIZE)
    draw = ImageDraw.Draw(image)
    for row_num in range(0, len(submission)):
        for column_num in range(0, len(submission[row_num])):
            x0 = column_num * BOX_SIZE
            y0 = row_num * BOX_SIZE
            x1 = x0 + BOX_SIZE
            y1 = y0 + BOX_SIZE
            draw.rectangle(
                [x0, y0, x1, y1],
                fill=submission[row_num][column_num]
            )
    del draw
    image.save(filename, "PNG")

# Create heatmap
heatmap = []
max_value = 0
for heatmap_row in range(0, GRID_SIZE):
    heatmap.append([0] * GRID_SIZE)
for submission in submissions:
    for submission_row_number, submission_row in enumerate(submission):
        for submission_column_number, submission_column in enumerate(submission_row):
            heatmap[submission_row_number][submission_column_number] += submission_column
            if heatmap[submission_row_number][submission_column_number] > max_value:
                max_value = heatmap[submission_row_number][submission_column_number]

# Print heatmap
ratio = 255 / max_value
image = Image.new("RGB", IMAGE_SIZE)
draw = ImageDraw.Draw(image)
for row_num in range(0, len(heatmap)):
    for column_num in range(0, len(heatmap[row_num])):
        x0 = column_num * BOX_SIZE
        y0 = row_num * BOX_SIZE
        x1 = x0 + BOX_SIZE
        y1 = y0 + BOX_SIZE
        draw.rectangle(
            [x0, y0, x1, y1],
            fill=(int(heatmap[row_num][column_num] * ratio), 0, 0)
        )
del draw
image.save("heatmap.png", "PNG")

data.csv