main.py

"""
Open Hoop POV Image Builder

Description:
The Open Hoop POV Image Builder is an essential tool within the OpenHoop project, which is dedicated to creating
interactive images using the Persistence of Vision (POV) technique.
This script enables users to create image files and corresponding code for displaying POV images on the OpenHoop device.

Key Features:
- Generates Persistence of Vision (POV) instructions for the OpenHoop device.
- Allows users to specify grid size and color palette for customized POV Image creation.
- Provides color aliasing functionality for easy management of variable colors.
- Automatically generates C++ class files (.h and .cpp) with appropriate methods for setting colors.
- Offers a user-friendly interface with prompt dialogs for input and feedback.

Author: github.com/angelcamelot
Date: 2024-04-09
License: Open-source license.

Requirements:
- PIL (Python Imaging Library)
- prompt_toolkit (for user input)

For more information about the OpenHoop project, visit:
https://github.com/angelcamelot/OpenHoop

Usage:
Run the script and follow the prompts.
"""

import os
import math
import re
import datetime
from PIL import Image, ImageDraw
from prompt_toolkit.shortcuts import input_dialog, message_dialog


def find_closest_color(target_color, colors):
    """Find the closest color in a list of colors to the target color."""
    closest_color = min(colors, key=lambda color: math.sqrt(sum((x - y) ** 2 for x, y in zip(target_color, color))))
    return closest_color


def rgb_to_str(color):
    """Convert an RGB color to string format."""
    return '{{{}, {}, {}}}'.format(color[0], color[1], color[2])


def generate_files(class_name, pov_code, pov_image, color_key_list=None):
    """Generates output files."""
    # Create a directory for the class_name inside the 'outputs' folder if it doesn't exist
    header_code = ""
    source_code = ""
    output_dir = os.path.join("outputs", class_name)
    current_date = datetime.datetime.now().strftime("%Y-%m-%d")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

        # Write header file (.h)
        header_code = f"""
/**
 * @project OpenHoop
 * @file {class_name}.cpp
 * @brief Source file for the {class_name} class.
 * @details Implements methods to set colors for the POV image.
 * @author Automatically generated by github.com/angelcamelot/OpenHoopPOVImageBuilder
 * @date {current_date}
 * @license Open-source license.
 */

#ifndef OPENHOOP_{class_name.upper()}_H
#define OPENHOOP_{class_name.upper()}_H

#include "POVImage.h"

/**
 * @brief Represents an image of POV Image for the Hula Hoop LED display.
 */
class {class_name.capitalize()} : public POVImage {{
public:
    /**
     * @brief Constructor for the {class_name.capitalize()} class.
     */
    explicit {class_name.capitalize()}({
            ', '.join([
                f'LedColor {color_key.capitalize()}' for color_key in color_key_list
            ]) if color_key_list else ''
        });

private:
    /**
     * @brief Set the colors for the {class_name}.
     */
    void setColors({
            ', '.join([
                f'LedColor {color_key.capitalize()}' for color_key in color_key_list
            ]) if color_key_list else ''
        });
}};

#endif //OPENHOOP_{class_name.upper()}_H

        """

        # Write the source file (.h)
        source_code = f"""
/**
 * @project OpenHoop
 * @file {class_name}.cpp
 * @brief Source file for the {class_name} class.
 * @details Implements methods to set colors for the POV image.
 * @author Automatically generated by github.com/angelcamelot/OpenHoopPOVImageBuilder
 * @date {current_date}
 * @license Open-source license.
 */

#include "../../include/images/{class_name}.h"

{class_name}::{class_name}({
            ', '.join([
                f'LedColor {color_key.capitalize()}' for color_key in color_key_list
            ]) if color_key_list else ''
        }) : POVImage({pov_image.size[0]}, {pov_image.size[1]}) {{
    setColors({
            ', '.join([
                f'LedColor {color_key.capitalize()}' for color_key in color_key_list
            ]) if color_key_list else ''
        });
}}

void {class_name}::SetColors({
            ', '.join([
                f'LedColor {color_key.capitalize()}' for color_key in color_key_list
            ]) if color_key_list else ''
        }) {{
{pov_code}
}}

    """

    # Write the POV instructions code to .h and .cpp files
    with open(os.path.join(output_dir, class_name + ".h"), "w") as f:
        f.write(header_code)
    with open(os.path.join(output_dir, class_name + ".cpp"), "w") as f:
        f.write(source_code)

    # Save the new image file for preview...
    pov_image.save(os.path.join(output_dir, class_name + ".bmp"))


def generate_pov_from_bitmap(image_path, class_name):
    """Generate POV instructions from a bitmap image."""
    # Open the image
    original_image = Image.open(image_path)

    # Convert the image to RGB mode to ensure compatibility
    original_image = original_image.convert("RGB")

    # Get the image size
    width, height = original_image.size

    # Create a new blank image
    pov_image = Image.new("RGB", (width, height), color=(255, 255, 255))
    draw = ImageDraw.Draw(pov_image)

    pov_code = ""

    # Iterate over each pixel in the image
    for y in range(height):
        for x in range(width):
            # Get the color index of the current pixel
            pixel_color = original_image.getpixel((x, y))

            # Write the instruction to the file
            hex_color = rgb_to_str(pixel_color)
            pov_code += f"    setPixel({x}, {y}, {hex_color});\n"

            # Draw the pixel with the same color
            draw.point((x, y), fill=pixel_color)

    generate_files(class_name, pov_code, pov_image)


def generate_pov_from_image(image_path, class_name, size, color_palette, color_key_mapping):
    """Generate POV instructions from a given image and a list of colors."""
    # Open the image
    original_image = Image.open(image_path)

    # Get the image size
    width, height = original_image.size

    # Calculate square size
    square_width = int(width) // int(size[0])
    square_height = int(height) // int(size[1])

    # Create a new blank image
    pov_image = Image.new("RGB", (int(size[0]), int(size[1])), color=0)
    draw = ImageDraw.Draw(pov_image)

    # POV code
    pov_code = ""
    color_key_list = []
    # Iterate over each square in the grid
    for y in range(int(size[1])):
        for x in range(int(size[0])):
            # Calculate the position of the current square
            x0 = x * square_width
            y0 = y * square_height

            # Get the color of the central pixel of the current square in the original image
            pixel_color = original_image.getpixel((x0 + square_width // 2, y0 + square_height // 2))

            # Find the closest color in the color list
            closest_color_rgb = find_closest_color(pixel_color, color_palette)

            # Get the color alias if present
            alias = color_key_mapping.get(closest_color_rgb)

            # If there is an alias, and it's not for background, write the instruction to the file
            if alias and str(alias).strip("'").strip('"') != str('background'):
                if str(alias).strip("'").strip('"') not in color_key_list:
                    color_key_list.append(str(alias).strip("'").strip('"'))
                pov_code += f"    setPixel({x}, {y}, {str(alias).strip("'").strip('"')});\n"
            elif not alias:
                hex_color = rgb_to_str(closest_color_rgb)
                pov_code += f"    setPixel({x}, {y}, {hex_color});\n"

            # Draw the square with the closest color
            draw.rectangle([(x, y), (x + 1, y + 1)], fill=closest_color_rgb)

    generate_files(class_name, pov_code, pov_image, color_key_list)


def split_commas(s):
    """Split a string by commas and return a list of parts."""
    parts = re.split(r',(?![^\(]*\))', s)  # noqa
    return [part.strip() for part in parts if part.strip()]


if __name__ == "__main__":
    message_dialog(
        title="OpenHoop POV Image Builder",
        text='Press ENTER to start.',
        ok_text='Start'
    ).run()
    while True:
        image_path_input = input_dialog(
            title="OpenHoop POV Image Builder",
            text="Please enter the path to the image file (supported file formats: .bmp, .png, .jpg, .jpeg):"
        ).run()

        if os.path.isfile(image_path_input) and image_path_input.lower().endswith(('.bmp', '.png', '.jpg', '.jpeg')):
            break
        else:
            message_dialog(
                title="Validation Error",
                text="Invalid file path or unsupported file format. Please enter a valid image file path.",
                ok_text="Retry"
            ).run()

    while True:
        class_name_input = input_dialog(
            title="OpenHoop POV Image Builder",
            text="Enter the name for the output class (maximum 20 characters, starting with a capital letter):"
        ).run()

        if re.match(r"^[A-Z][a-zA-Z0-9]{0,19}$", class_name_input):
            break
        else:
            message_dialog(
                title="Validation Error",
                text="Invalid class name.",
                ok_text="Retry"
            ).run()
    if image_path_input.lower().endswith('.bmp'):
        generate_pov_from_bitmap(image_path_input, class_name_input)
    else:
        while True:
            size_input = input_dialog(
                title="OpenHoop POV Image Builder",
                text="Please enter the grid size (width, height), separated by comma, e.g., '30,30':"
            ).run()

            size_input = tuple(map(int, split_commas(size_input)))

            if len(size_input) == 2 and all(isinstance(x, int) for x in size_input) and all(x > 0 for x in size_input):
                break
            else:
                message_dialog(
                    title="Validation Error",
                    text="Invalid grid size. Please enter two positive integers separated by comma.",
                    ok_text="Retry"
                ).run()

        while True:
            color_palette_input = input_dialog(
                title="OpenHoop POV Image Builder",
                text="Enter the list of colors in RGB format, separated by comma, e.g., '(52, 32, 20), (164, 102, 55)':"
            ).run()

            try:
                color_palette_input = [eval(color) for color in split_commas(color_palette_input)]
                if all(isinstance(color, tuple) and len(color) == 3 and all(0 <= x <= 255 for x in color) for color in
                       color_palette_input):
                    break
                else:
                    raise ValueError
            except (ValueError, SyntaxError):
                message_dialog(
                    title="Validation Error",
                    text="Invalid color palette. Please enter a list of RGB tuples.",
                    ok_text="Retry"
                ).run()

        while True:
            color_key_mapping_input = input_dialog(
                title="OpenHoop POV Image Builder",
                text="Enter color aliases in key:value format, separated by comma,"
                     "e.g., '(52, 32, 20):bkg, (164, 102, 55):primary_color, (206, 143, 81):secondary_color':"
            ).run()

            color_key_mapping_input = split_commas(color_key_mapping_input)
            color_key_mapping_list = {}

            for item in color_key_mapping_input:
                key_value_pair = item.split(':')
                try:
                    color = eval(key_value_pair[0])
                    if len(color) != 3 or not all(0 <= x <= 255 for x in color):
                        raise ValueError
                    color_key_mapping_list[color] = key_value_pair[1]
                except (ValueError, SyntaxError):
                    message_dialog(
                        title="Validation Error",
                        text="Invalid color format in color alias. Please enter a valid RGB tuple.",
                        ok_text="Retry"
                    ).run()
                    break
            else:
                break

        color_key_mapping_input = color_key_mapping_list
        generate_pov_from_image(
            image_path_input, class_name_input, size_input, color_palette_input, color_key_mapping_input
        )
    message_dialog(
        title="OpenHoop POV Image Builder",
        text="Thank you for using OpenHoop POV Image Builder. Press ENTER to exit."
    ).run()