sierpinski_complete.py

from math import cos, pi, sin
import cv2
import numpy
from random import choice, randint
from PIL import Image
import click


WHITE = (255, 255, 255)
RED = (0, 0, 255)
BLACK = (0, 0, 0)


def draw_circle(image, r, padding):
    for i in range(2000):
        point_x = int(r * cos(2 * pi * (i+1) / 2000))
        point_y = int(r * sin(2 * pi * (i+1) / 2000))
        image[r + padding + point_x][r + padding + point_y] = RED


# Returns the midway point between the given point and one of the vertexes, picked randomly
def get_point(point, initial_points):
    p = choice(initial_points)
    return (int((point[0] + p[0]) / 2), int((point[1] + p[1]) / 2))


@click.command()
@click.option('--iterations', default=1_000_000, help='Number of iterations to perform.')
@click.option('--sides', default=3, help='Sides of the polygon to generate.')
@click.option('--size', default=1000, help='Size of the image (without padding).')
@click.option('--padding', default=50, help='White border around the image.')
@click.option('--image-rate', default=1_000, help='Generate a new frame of the video every X iterations.')
def generate(iterations, sides, size, padding, image_rate):
    r = size // 2
    image_size = size + padding * 2

    output_prefix = f"chaos_generator_{sides}_sided_{iterations}_iterations"
    video_writer = cv2.VideoWriter(f'{output_prefix}.avi', cv2.VideoWriter_fourcc('M','J','P','G'), 60, (image_size, image_size))

    # Create the base image of the right size making it all white. Every point is a tuple (R, G, B)
    image = numpy.ones((image_size, image_size, 3), dtype=numpy.uint8) * 255
    # Draw the circumscribed circle
    draw_circle(image, r, padding)

    initial_points = []
    # Generate the initial points and draw them
    for i in range(sides):
        point_x = int(r * cos(2 * pi * (i+1) / sides))
        point_y = int(r * sin(2 * pi * (i+1) / sides))
        # Translate the points for the image canvas
        initial_points.append((r + padding + point_x, r + padding + point_y))
        image[point_x][point_y] = BLACK

    # Randomly pick the start point
    current_point = (padding + randint(1, size - 1), padding + randint(1, size - 1))
    image[current_point[0]][current_point[1]] = BLACK
    video_writer.write(image)

    # Add a point at every iteration
    for i in range(1, iterations):
        current_point = get_point(current_point, initial_points)
        image[current_point[0]][current_point[1]] = BLACK
        if i % image_rate == 0:
            video_writer.write(image)

    video_writer.release()
    image = Image.fromarray(image)
    image.save(f"{output_prefix}.png")


if __name__ == '__main__':
    generate()