draw.py

import drawsvg as dw
import cairosvg
from io import StringIO
import numpy as np

def generate_grid_image(grid, output_filename="grid_image.png", cell_size=20, bg_color="#f0f0f0", padding=10, extra_bottom_padding_ratio=0.5, add_size=True, label="", bordercol="black"):
    if not isinstance(grid, np.ndarray):
        grid = np.array(grid)
    grid = grid.astype(int)
        
    cmap = [
        '#252525', '#0074D9', '#FF4136', '#37D449', '#FFDC00',
        '#808080', '#F012BE', '#FF871E', '#54D2EB', '#8D1D2C', '#FFFFFF'
    ]
    
    gridy, gridx = grid.shape
    
    grid_pixel_width = gridx * cell_size
    grid_pixel_height = gridy * cell_size
    extra_bottom_padding_px = cell_size * extra_bottom_padding_ratio

    total_image_width = grid_pixel_width + 2 * padding
    total_image_height = grid_pixel_height + 2 * padding + extra_bottom_padding_px

    line_thickness = 0.5
    circle_radius = 0
    border_width = 1.0
    lt = line_thickness / 2

    drawing = dw.Drawing(total_image_width, total_image_height, origin=(-padding, -padding))
    
    drawing.append(dw.Rectangle(
        -padding, -padding,
        total_image_width, total_image_height,
        fill=bg_color
    ))
    
    for j, row in enumerate(grid):
        for i, cell in enumerate(row):
            cell_index = int(cell) % len(cmap)
            drawing.append(dw.Rectangle(i*cell_size+lt, j*cell_size+lt, cell_size-lt, cell_size-lt, fill=cmap[cell_index]))
    
    if circle_radius > 0:
        for i in range(1, gridx):
            for j in range(1, gridy):
                drawing.append(dw.Circle(i*cell_size, j*cell_size, circle_radius, fill='white'))
        
    drawing.append(dw.Rectangle(-padding, -padding, total_image_width, total_image_height, fill='none', stroke=bordercol, stroke_width=border_width))

    
    
    svg_string = drawing.as_svg()
    cairosvg.svg2png(bytestring=svg_string.encode('utf-8'), write_to=output_filename, dpi=150)