Image to Gaussian Mixture

import numpy as np
from skimage import io, color
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', 'lab', or 'grey').

    Returns:
    - data (numpy.ndarray): The prepared data containing normalized coordinates and either RGB, Lab, 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 == 'lab':
        # Convert RGB to Lab color space
        lab_values = color.rgb2lab(image_resized).reshape(-1, 3)

        # Normalize to [0,1] range
        lab_values[:, 0] = lab_values[:, 0] / 100.0  # Scale L to [0,1]
        lab_values[:, 1] = (lab_values[:, 1] + 128) / 255.0  # Scale a to [0,1]
        lab_values[:, 2] = (lab_values[:, 2] + 128) / 255.0  # Scale b to [0,1]

        values = lab_values
    elif mode == 'grey':
        # Luminosity method: Y = 0.299 * R + 0.587 * G + 0.114 * B
        values = np.dot(image_resized[..., :3], [0.299, 0.587, 0.114]).flatten().reshape(-1, 1)
    else:
        raise ValueError("Invalid mode. Use 'rgb', 'lab', or 'grey'.")

    # 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  # RGB, Lab, or grayscale values
    ))

    return data


def prepare_image_data_with_mask(image_path, mask_path, num_samples, mode='rgb', random_seed=None):
    """
    Loads an image and a focus mask, then selects pixels randomly based on the mask brightness.

    Args:
    - image_path (str): Path to the input image.
    - mask_path (str): Path to the grayscale focus mask. The brighter, the more important.
    - num_samples (int): Total number of pixels to sample.
    - mode (str): Mode for processing the image data ('rgb', 'lab', or 'grey').
    - random_seed (int, optional): Seed for reproducible sampling. If None, results are random.

    Returns:
    - data (numpy.ndarray): The selected data containing coordinates, color values, and focus weights.
    """
    # Load images
    image = io.imread(image_path) / 255.0  # Normalize image to [0,1]
    mask = io.imread(mask_path) / 255.0

    # If the mask has multiple channels, take only the first (red channel)
    if mask.ndim == 3:
        mask = mask[..., 0]  # Select the first channel (Red)

    # Get image dimensions
    height, width = image.shape[:2]

    # Create coordinate grid
    y_coords, x_coords = np.meshgrid(np.linspace(1, 0, height), np.linspace(0, 1, width), indexing='ij')

    # Flatten data
    x_coords = x_coords.flatten()
    y_coords = y_coords.flatten()
    mask_values = mask.flatten()

    if mode == 'rgb':
        values = image.reshape(-1, 3)  # RGB values
    elif mode == 'lab':
        # Convert RGB to Lab color space
        lab_values = color.rgb2lab(image).reshape(-1, 3)

        # Normalize to [0,1] range
        lab_values[:, 0] = lab_values[:, 0] / 100.0  # Scale L to [0,1]
        lab_values[:, 1] = (lab_values[:, 1] + 128) / 255.0  # Scale a to [0,1]
        lab_values[:, 2] = (lab_values[:, 2] + 128) / 255.0  # Scale b to [0,1]

        values = lab_values
    elif mode == 'grey':
        # Luminosity method: Y = 0.299 * R + 0.587 * G + 0.114 * B
        values = np.dot(image[..., :3], [0.299, 0.587, 0.114]).flatten().reshape(-1, 1)
    else:
        raise ValueError("Invalid mode. Use 'rgb' or 'grey'.")

    # Normalize mask to use as probabilities
    mask_values += 1e-6  # Avoid division by zero
    mask_values /= np.sum(mask_values)

    # Set random seed if provided
    if random_seed is not None:
        np.random.seed(random_seed)

    # Sample pixels based on mask probabilities
    indices = np.random.choice(len(mask_values), size=num_samples, p=mask_values)

    # Select data points
    sampled_x = x_coords[indices]
    sampled_y = y_coords[indices]
    sampled_values = values[indices]

    # Combine into final data array
    data = np.column_stack((sampled_x, sampled_y, sampled_values))

    return data


def find_scaling_factor(values, target_digits=3):
    """
    Finds the scaling factor for a given dataset to ensure values are represented as
    integers with the specified number of digits. This scaling factor will normalize
    the data such that the largest value is scaled to a 3-digit integer.

    Args:
        values (numpy.ndarray): The dataset to find the scaling factor for.
        target_digits (int, optional): The number of digits to scale the values to.
                                       Defaults to 3.

    Returns:
        float: The scaling factor to scale the values to the target number of digits.
    """
    # Find the largest value in the dataset
    max_val = np.max(np.abs(values))
    if max_val == 0:
        return 1  # Avoid division by zero if values are zero

    # Find the scaling factor (targeting a 3-digit integer)
    scale_factor = 10 ** (np.floor(np.log10(max_val)) - (target_digits - 1))
    return scale_factor


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

    Args:
        gmm (sklearn.mixture.GaussianMixture): The trained Gaussian Mixture Model object containing:
            - `weights_`: The mixture component weights.
            - `means_`: The cluster means (centroids) of the GMM.
            - `covariances_`: The covariance matrices of the GMM.

    Returns:
        str: A GLSL-compatible string that includes:
            - `COUNT`: The number of Gaussian components.
            - `SCALES`: A `vec4` containing scaling factors for weights, positions, colors, and covariances.
            - `WEIGHTS`: A GLSL `int` array representing the scaled weights of the components.
            - `COLORS`: A GLSL `vec3` array representing the scaled RGB or brightness values.
            - `POSITIONS`: A GLSL `vec2` array representing the scaled positions (x, y) of the components.
            - `COVARIANCES`: A GLSL `mat2` array representing the scaled 2x2 covariance matrices.
    """
    # Extracting cluster means and covariances from the GMM
    weights = gmm.weights_
    cluster_means = gmm.means_
    covariances = gmm.covariances_

    # Inferred number of components
    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

    # Extract the first 4 covariance elements affecting the position only
    cov_selected = np.array([cov[:2, :2].flatten() for cov in covariances])

    # Compute scaling factors for each type of value
    weight_scale = find_scaling_factor(weights)
    position_scale = find_scaling_factor(cluster_means[:, :2])
    if use_rgb:
        color_scale = find_scaling_factor(cluster_means[:, 2:5])
    else:
        color_scale = find_scaling_factor(cluster_means[:, 2:3])
    cov_scale = find_scaling_factor(cov_selected, target_digits=4)

    # Prepare GLSL output
    output = f"\n#define COUNT {n_components}\n"
    output += f"#define SCALES vec4({weight_scale:.1e},{position_scale:.1e},{color_scale:.1e},{cov_scale:.1e})\n"
    output = output.replace("1.0e-0", "1e-")  # Safe some characters

    # Weights (scaled to 3 digits)
    output += "const int WEIGHTS[COUNT] = int[](" + ",".join(f"{int(w / weight_scale):d}" for w in weights) + ");\n"

    # Colors (scaled to 3 digits)
    output += "const vec3 COLORS[COUNT] = vec3[]("
    output += ",".join(f"v({int(r / color_scale):d},{int(g / color_scale):d},{int(b / color_scale):d})"
                       if use_rgb else f"vec3({int(r / color_scale):d})"
                       for r, g, b in (cluster_means[:, 2:5]
                                       if use_rgb else np.c_[cluster_means[:, 2], cluster_means[:, 2], cluster_means[:, 2]]))
    output += ");\n"

    # Positions (scaled to 3 digits)
    output += "const vec2 POSITIONS[COUNT] = vec2[]("
    output += ",".join(f"u({int(x / position_scale):d},{int(y / position_scale):d})"
                       for x, y in cluster_means[:, :2])
    output += ");\n"

    # Covariances (2x2 sampled, scaled to 3 digits)
    output += "const mat2 COVARIANCES[COUNT] = mat2[]("
    output += ",".join(f"m({int(c[0] / cov_scale):d},{int(c[1] / cov_scale):d},{int(c[2] / cov_scale):d},{int(c[3] / cov_scale):d})"
                       for c in cov_selected)
    output += ");"

    return output


# Path to image to convert and down-sampling and number of clusters
# example images can be found here https://imgur.com/a/zQYTfFQ
image_path = 'Textures/Lenna.png'
mask_path = 'Textures/Lenna_Importance.png'  # optional importance mask made with MS Paint
new_width = 128  # increase for better quality results
clusters = 512

# data = prepare_image_data(image_path, new_width, mode='lab')
data = prepare_image_data_with_mask(image_path, mask_path, new_width * new_width, mode='lab')

# Fit the Gaussian Mixture Model
model = GaussianMixture(n_components=clusters, max_iter=100, covariance_type='full', init_params='k-means++', verbose=2)

# Claims it's a better model but takes 3x as long and makes very blurry results
# model = BayesianGaussianMixture(n_components=clusters, max_iter=300, covariance_type='full', init_params='k-means++', verbose=2)

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