visualization/draw.py

from pathlib import Path

from PIL import Image, ImageDraw

from .corners import Circle, Corners

DEFAULT_COLOR = (0, 255, 255, 255)


def draw_circle(draw: ImageDraw.Draw, circle: Circle, color: tuple, width: int = 3):
    x1 = circle.x - circle.r
    y1 = circle.y - circle.r
    x2 = circle.x + circle.r
    y2 = circle.y + circle.r
    draw.ellipse([x1, y1, x2, y2], outline=color, width=width)


def draw_circles(
    image: Image.Image,
    corners: Corners,
    color: tuple = DEFAULT_COLOR,
    width: int = 3,
) -> Image.Image:
    image = image.convert("RGBA")
    overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    for corner in [corners.top_left, corners.top_right, corners.bottom_left, corners.bottom_right]:
        draw_circle(draw, corner, color, width)

    result = Image.alpha_composite(image, overlay)
    return result.convert("RGB")


def draw_crop_lines(
    image: Image.Image,
    corners: Corners,
    color: tuple = DEFAULT_COLOR,
    width: int = 2,
) -> Image.Image:
    image = image.convert("RGBA")
    overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    x_left, y_top, x_right, y_bottom = corners.get_crop_bounds()
    img_w, img_h = image.size

    draw.line([(x_left, 0), (x_left, img_h)], fill=color, width=width)
    draw.line([(x_right, 0), (x_right, img_h)], fill=color, width=width)
    draw.line([(0, y_top), (img_w, y_top)], fill=color, width=width)
    draw.line([(0, y_bottom), (img_w, y_bottom)], fill=color, width=width)

    result = Image.alpha_composite(image, overlay)
    return result.convert("RGB")


def visualize_corners(
    image: Image.Image,
    corners: Corners,
    output_path: Path | None = None,
    color: tuple = DEFAULT_COLOR,
) -> Image.Image:
    image = image.convert("RGBA")
    overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    for corner in [corners.top_left, corners.top_right, corners.bottom_left, corners.bottom_right]:
        draw_circle(draw, corner, color, width=3)

    x_left, y_top, x_right, y_bottom = corners.get_crop_bounds()
    img_w, img_h = image.size

    draw.line([(x_left, 0), (x_left, img_h)], fill=color, width=2)
    draw.line([(x_right, 0), (x_right, img_h)], fill=color, width=2)
    draw.line([(0, y_top), (img_w, y_top)], fill=color, width=2)
    draw.line([(0, y_bottom), (img_w, y_bottom)], fill=color, width=2)

    result = Image.alpha_composite(image, overlay)
    result = result.convert("RGB")

    if output_path:
        result.save(output_path)

    return result