Lenna Gaussian Mixture

import numpy as np
from skimage import io
from skimage.transform import resize
from sklearn.mixture import GaussianMixture
from sklearn.mixture import BayesianGaussianMixture


def prepare_image_data(image_path, new_width, mode='rgb'):
    """
    Loads an image, resizes it, and prepares the data for clustering.

    Args:
    - image_path (str): Path to the input image.
    - new_width (int): New width for the resized image, aspect ratio will be preserved.
    - mode (str): Mode for processing the image data ('rgb' or 'brightness').

    Returns:
    - data (numpy.ndarray): The prepared data containing normalized coordinates and either RGB values or brightness values.
    """
    # Load the image
    image = io.imread(image_path)

    # Get original dimensions
    orig_height, orig_width, _ = image.shape

    # Set new height to preserve aspect ratio
    new_height = int((new_width / orig_width) * orig_height)

    # Resize the image while preserving aspect ratio
    image_resized = resize(image, (new_height, new_width), anti_aliasing=True)

    if mode == 'rgb':
        # Extract RGB values and normalize them to [0, 1]
        values = image_resized.reshape(-1, 3)
    elif mode == 'brightness':
        # Calculate brightness as max(r, g, b) and normalize to [0, 1]
        values = np.max(image_resized, axis=2).flatten().reshape(-1, 1)

    # Create coordinate grid
    y_indices, x_indices = np.meshgrid(
        np.linspace(1, 0, new_height), np.linspace(0, 1, new_width), indexing='ij'
    )

    # Reorder data so that coordinates come first (x, y, values)
    data = np.column_stack((
        x_indices.flatten(),  # Normalized x-coordinates
        y_indices.flatten(),  # Normalized y-coordinates
        values  # Flattened values (either RGB or brightness)
    ))

    return data


def generate_glsl_data(gmm):
    """
    Converts clustering results from a Gaussian Mixture Model into GLSL-ready data definitions.

    Args:
    - gmm (GaussianMixture): The fitted GMM object.

    Returns:
    - str: GLSL code string containing the cluster information.
    """
    # Extracting cluster means and covariances from the GMM
    weights = gmm.weights_
    cluster_means = gmm.means_
    covariances = gmm.covariances_
    covariance_type = gmm.covariance_type  # Get the covariance type

    # Inferred n_components from the length of cluster_means
    n_components = len(cluster_means)

    # Determine if using RGB or brightness mode
    use_rgb = cluster_means.shape[1] == 5  # Assuming [x, y, r, g, b] for RGB mode

    # Prepare GLSL output
    output = f"#define COUNT {n_components}\n\n"
    # Example Python code to generate GLSL code for weights

    # Printing weight of a cluster definitions
    output += "    const float weights[COUNT] = float[]("
    for i, weight in enumerate(weights):
        output += f"{weight:.2e}"
        if i < len(weights) - 1:
            output += ","
    output += ");\n"

    # Printing color definitions
    output += "    const vec3 colors[COUNT] = vec3[]("
    for i, mean in enumerate(cluster_means):
        if use_rgb:
            r, g, b = mean[2], mean[3], mean[4]
            output += f"vec3({r:.3f},{g:.3f},{b:.3f})"
        else:
            r = mean[2]  # Use the third value for grayscale
            output += f"vec3({r:.3f})"
        if i < len(cluster_means) - 1:
            output += ","
    output += ");\n"

    # Printing coordinate center definitions
    output += "    const vec2 positions[COUNT] = vec2[]("
    for i, mean in enumerate(cluster_means):
        x, y = mean[0], mean[1]
        output += f"vec2({x:.3f},{y:.3f})"
        if i < len(cluster_means) - 1:
            output += ","
    output += ");\n"

    # Handle different covariance types
    output += "    const mat2 covariances[COUNT] = mat2[]("
    for i, covariance in enumerate(covariances):
        if covariance_type == 'spherical':
            variance = covariance  # Spherical covariance is a single variance value
            output += f"mat2({variance:.2e},0,0,{variance:.2e})"
        elif covariance_type == 'full':
            c00, c01 = covariance[0, 0], covariance[0, 1]
            c10, c11 = covariance[1, 0], covariance[1, 1]
            output += f"mat2({c00:.2e},{c01:.2e},{c10:.2e},{c11:.2e})"
        if i < len(covariances) - 1:
            output += ","
    output += ");\n"

    # Replace "0." with "."
    output = output.replace("0.", ".")

    return output


# Path to image to convert and down-sampling and number of clusters
image_path = 'Textures/Lenna.png'
new_width = 128
clusters = 512

data = prepare_image_data(image_path, new_width, mode='rgb')

# Fit the Gaussian Mixture Model
model = GaussianMixture(n_components=clusters, covariance_type='full', verbose=2)

# Fit the Bayesian Gaussian Mixture Model (claims it's better but looks blurry and is 2x slower)
# model = BayesianGaussianMixture(n_components=clusters, covariance_type='full', verbose=2)

# Run the model
model.fit(data)

# Print the GLSL code to use in your shader
print(generate_glsl_data(model))

# See Shader code at https://www.shadertoy.com/view/43Gfzt