turtle.py

from basalt import *
from typing import List

width = 1000
height = Variable()
border_radius = 10
grid_spacing = 10

# some geometric primitives that would be part of a library

def inset(inner, outer, padding_x=None, padding_y=None):
    padding_x = maybe_create(padding_x)
    if padding_y is None:
        padding_y = padding_x
    return Set([
        padding_x >= 0,
        padding_y >= 0,
        inner.bounds.center == outer.bounds.center,
        inner.bounds.left == outer.bounds.left + padding_x,
        inner.bounds.right == outer.bounds.right - padding_x,
        inner.bounds.top == outer.bounds.top + padding_y,
        inner.bounds.bottom == outer.bounds.bottom - padding_y,
    ])

def same_size(elements):
    width = Variable()
    height = Variable()
    return Set([And(i.bounds.width == width, i.bounds.height == height) for i in elements])

def distributed_horizontally(l, spacing=None):
    spacing = maybe_create(spacing)
    constrs = [spacing >= 0]
    last = l[0]
    for elem in l[1:]:
        constrs.append(last.bounds.right + spacing == elem.bounds.left)
        last = elem
    return Set(constrs)

def distributed_vertically(l, spacing=None):
    spacing = maybe_create(spacing)
    constrs = [spacing >= 0]
    last = l[0]
    for elem in l[1:]:
        constrs.append(last.bounds.bottom + spacing == elem.bounds.top)
        last = elem
    return Set(constrs)

def aligned_middle(l):
    middle = Variable()
    return Set([elem.bounds.center.y == middle for elem in l])

def aligned_center(l):
    center = Variable()
    return Set([elem.bounds.center.x == center for elem in l])

# object abstractions

def bordered_image(source: str, border_color, border_radius):
    image = Image(source)
    border = Rectangle(style=Style({Property.fill: border_color}))
    return Group([border, image], [
        inset(image, border, border_radius, border_radius)
    ])

def grid(elements: List[List[Element]], spacing):
    rows = [Group(i) for i in elements]
    return Group(rows, [
        # all elements are the same size
        same_size([i for row in rows for i in row.elements]),
        # elements in rows are distributed horizontally
        Set([distributed_horizontally(row, spacing) for row in rows]),
        # and aligned vertically
        Set([aligned_middle(row) for row in rows]),
        # rows are distributed vertically
        Set([distributed_vertically(rows, spacing)]),
        # and aligned horizontally
        aligned_center(rows)
    ])

colors = {'adversarial': RGB(0xd32f2f), 'correct': RGB(0x27ae60), 'misclassified': RGB(0x000000)}

across = 6
down = 4

bg = Rectangle(0, 0, width, height, style=Style({Property.fill: RGB(0xeeeeee)}))

images = [[None for _ in range(across)] for _ in range(down)]
for i in range(across*down):
    prefix = './turtle/{:04d}'.format(i)
    image_source = prefix + '.png'
    color = colors[open(prefix + '.txt').read().strip()]
    images[i//across][i%across] = bordered_image(image_source, color, border_radius)
grid = grid(images, grid_spacing)

group = Group(
    [bg, grid],
    [inset(grid, bg, 0)]
)

c = Canvas(group, width=width, height=height)