hubduck.py

import PIL
import PIL.Image
import PIL.ImageDraw
import numpy as np
import tqdm

image_in = "image_in.png"


def sample_pixel(image: PIL.Image.Image, x, y):
    """
    Sample a pixel from an image, with a cost function.

    If `x` and `y` are integers, this function returns the value of the pixel at
    that location. If they are floats, it returns the maximum of the four pixels
    surrounding the point.
    """

    if x < 0 or y < 0 or x >= image.width or y >= image.height:
        return None

    if x == int(x) and y == int(y):
        return image.getpixel((x, y))

    x0 = int(x)
    x1 = x0 + 1
    y0 = int(y)
    y1 = y0 + 1

    return max(
        image.getpixel((x0, y0)),
        image.getpixel((x0, y1)),
        image.getpixel((x1, y0)),
        image.getpixel((x1, y1)),
    )


def integrate_along(image: PIL.Image.Image, start, end, color_to_value):
    """
    Integrate along a line from start to end, with a cost function.

    This function is actually a numeric approximation of the integral by sampling
    the image along the line at 0.5px intervals and summing the results.
    """

    x0, y0 = start
    x1, y1 = end

    dx = x1 - x0
    dy = y1 - y0

    length = np.sqrt(dx * dx + dy * dy)

    if length == 0:
        return 0

    dx /= length
    dy /= length

    total = 0

    for i in range(int(length * 2)):
        x = x0 + dx * i / 2
        y = y0 + dy * i / 2

        value = sample_pixel(image, x, y)

        if value is None:
            continue

        total += color_to_value(value)

    return total


def pick_next_line(image: PIL.Image.Image, n: int):
    """
    Pick the next line to draw.

    - Generate `n` random lines, whose start and end points are also random within
        the image.
    - For each line, calculate the integral of darkness (1-brightness) along the
        line.
    - Return the line with the highest integral.
    """

    best_integral = 0
    best_line = None

    for _ in range(n):
        x0 = np.random.uniform(0, image.width - 1)
        y0 = np.random.uniform(0, image.height - 1)
        x1 = np.random.uniform(0, image.width - 1)
        y1 = np.random.uniform(0, image.height - 1)

        integral = integrate_along(
            image,
            (x0, y0),
            (x1, y1),
            lambda color: 1 - color / 255,
        )
        # Penaltize lines that are too long
        integral /= (x1 - x0)**2 + (y1 - y0)**2

        if integral > best_integral:
            best_integral = integral
            best_line = (x0, y0), (x1, y1)

    return best_line


lines = []
image = PIL.Image.open(image_in)
image = image.convert("L")
for _ in tqdm.tqdm(range(4000), "find line"):
    line = pick_next_line(image, 10)
    lines.append(line)

image_out = PIL.Image.new("RGB", image.size, "white")
draw = PIL.ImageDraw.Draw(image_out)
for line in tqdm.tqdm(lines, "draw lines"):
    draw.line(line, fill="black", width=1)

image_out.save("image_out.png")