serpdotai/sklearn_vector_quantization

1import gradio as gr2import matplotlib.pyplot as plt3from sklearn.preprocessing import KBinsDiscretizer4import numpy as np5from typing import Tuple6 7 8def build_init_plot(img_array: np.ndarray) -> Tuple[str, plt.Figure]:9    init_text = (f"The dimension of the image is {img_array.shape}\n"10          f"The data used to encode the image is of type {img_array.dtype}\n"11          f"The number of bytes taken in RAM is {img_array.nbytes}")12 13    fig, ax = plt.subplots(ncols=2, figsize=(12, 4))14 15    ax[0].imshow(img_array, cmap=plt.cm.gray)16    ax[0].axis("off")17    ax[0].set_title("Rendering of the image")18    ax[1].hist(img_array.ravel(), bins=256)19    ax[1].set_xlabel("Pixel value")20    ax[1].set_ylabel("Count of pixels")21    ax[1].set_title("Distribution of the pixel values")22    _ = fig.suptitle("Original image")23 24    return init_text, fig25 26 27def build_compressed_plot(compressed_image, img_array, sampling: str) -> plt.Figure:28    compressed_text = (f"The number of bytes taken in RAM is {compressed_image.nbytes}\n"29          f"Compression ratio: {compressed_image.nbytes / img_array.nbytes}\n"30          f"Type of the compressed image: {compressed_image.dtype}")31    32    sampling = sampling if sampling == "uniform" else "K-Means"33    34    fig, ax = plt.subplots(ncols=2, figsize=(12, 4))35    ax[0].imshow(compressed_image, cmap=plt.cm.gray)36    ax[0].axis("off")37    ax[0].set_title("Rendering of the image")38    ax[1].hist(compressed_image.ravel(), bins=256)39    ax[1].set_xlabel("Pixel value")40    ax[1].set_ylabel("Count of pixels")41    ax[1].set_title("Sub-sampled distribution of the pixel values")42    _ = fig.suptitle(f"Original compressed using 3 bits and a {sampling} strategy")43 44    return compressed_text, fig45 46 47def infer(img_array: np.ndarray, sampling: str, number_of_bins: int):48    # greyscale_image = input_image.convert("L")49    # img_array = np.array(greyscale_image)50 51    #raccoon_face = face(gray=True)52    init_text, init_fig = build_init_plot(img_array)53    54    n_bins = number_of_bins55    encoder = KBinsDiscretizer(56        n_bins=n_bins, encode="ordinal", strategy=sampling, random_state=057    )58    compressed_image = encoder.fit_transform(img_array.reshape(-1, 1)).reshape(59        img_array.shape60    )61    compressed_image = compressed_image.astype(np.uint8)62    compressed_text, compressed_fig = build_compressed_plot(compressed_image,63                                                            img_array,64                                                            sampling)65 66    bin_edges = encoder.bin_edges_[0]67    bin_center = bin_edges[:-1] + (bin_edges[1:] - bin_edges[:-1]) / 268 69    comparison_fig, ax = plt.subplots()70    ax.hist(img_array.ravel(), bins=256)71    color = "tab:orange"72    for center in bin_center:73        ax.axvline(center, color=color)74        ax.text(center - 10, ax.get_ybound()[1] + 100, f"{center:.1f}", color=color)75 76    return init_text, init_fig, compressed_text, compressed_fig, comparison_fig77 78 79article = """<center>80Demo by <a href='https://huggingface.co/johko' target='_blank'>Johannes (johko) Kolbe</a>"""81 82 83gr.Interface(84    title="Vector Quantization with scikit-learn",85    description="""<p style="text-align: center;">This is an interactive demo for the <a href="https://scikit-learn.org/stable/auto_examples/cluster/plot_face_compress.html">Vector Quantization Tutorial</a> from scikit-learn.86    </br><b>Vector Quantization</b> is a compression technique to reduce the number of color values that are used in an image and with this save memory while trying to keep a good quality.87    In this demo this can be done naively via <i>uniform</i> sampling, which just uses <i>N</i> color values (specified via slider) uniformly sampled from the whole spectrum or via <i>k-means</i> which pays closer attention 88    to the actual pixel distribution and potentially leads to a better quality of the compressed image.89    In this demo we actually won't see a compression effect, because we cannot go smaller than <i>uint8</i> in datatype size here.90    </br>91    </br><b>Usage</b>: To run the demo you can simply upload an image and choose from two sampling methods - <i>uniform</i> and <i>kmeans</i>. Choose the number of bins and then click 'submit'. 92    You will get information about the histogram, pixels distribution and other image statistics for your orginial image as grayscale and the quantized version of it.93    </p>""",94    article=article,95    fn=infer, 96    inputs=[gr.Image(image_mode="L", label="Input Image"),97            gr.Dropdown(choices=["uniform", "kmeans"], label="Sampling Method"),98            gr.Slider(minimum=2, maximum=50, value=8, step=1, label="Number of Bins")], 99    outputs=[gr.Text(label="Original Image Stats"),100             gr.Plot(label="Original Image Histogram"),101             gr.Text(label="Compressed Image Stats"), 102             gr.Plot(label="Compressed Image Histogram"),103             gr.Plot(label="Pixel Distribution Comparison")],104    examples=[["examples/hamster.jpeg", "uniform", 8],105              ["examples/racoon.png", "kmeans", 8]]).launch()106