sklearn 实现极简VQ算法

VQ算法是一种图像压缩算法。本文用 sklearn 实现VQ算法。

代码简单的要哭

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from sklearn.mixture import *
from sklearn.cluster import *

from PIL import Image

def vq(im, model=None):
    # vector quantization with a clustering model
    # im: an image read by pillow
    # model: a clustering model
    if model is None:
        model = KMeans(n_clusters=2)

    I = np.asarray(im, dtype=np.float64)
    X = I.reshape((im.size[0]*im.size[1], 3))
 
    model.fit(X)
    code = model.predict(X)  # encoding
    if hasattr(model, 'cluster_centers_'):
        Y = np.array([model.cluster_centers_[k] for k in code])   # decode with the codebook 'model.cluster_centers_'
    elif hasattr(model, 'means_'):
        Y = np.array([model.means_[k] for k in code])   # decode with the codebook 'model.means_'
    Y = Y.reshape((im.size[1], im.size[0])+(3,))
    return Image.fromarray(Y.astype('uint8', copy=False))

def test():
    # 实例
    model = BayesianGaussianMixture(n_components=2)
    im = Image.open('lenna.jpg')
    im = vq(im, model)
    im.save('lenna.vq2.jpg')

注：下图采用的是Kmeans

面向对象风格

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""VQ algorithm for image compression

Reference: https://www.sciencedirect.com/topics/engineering/vector-quantization
"""

import numpy as np
from sklearn.mixture import *
from sklearn.cluster import *

from PIL import Image

class ImageTransformer:
    pass

class VQTransformer(ImageTransformer):
    """VQ algorithm for image compression

    Example:
        model = KMeans(n_clusters=2)
        im = Image.open('cloth.jpg')
        t = VQTransformer(model=model)
        t.fit(im)
        im = t.transform(im)
        im.save('cloth.kmeans2.jpg')
    """

    def __init__(self, small_size=(1, 1), model=None):
        super().__init__()
        if model is None:
            model = KMeans(n_clusters=2)
        self.small_size = small_size
        self.model = model

    def fit(self, im, shape=None):
        X, shape = image2data(im, small_size=self.small_size, shape=shape)
        self.shape = shape
        self.model.fit(X)
        if hasattr(self.model, 'cluster_centers_'):
            self.codebook_ = self.model.cluster_centers_
        elif hasattr(self.model, 'means_'):
            self.codebook_ = self.model.means_
        return self

    def get_code(self, k):
        w, h = self.small_size
        return self.codebook_[k].reshape((h, w, 3))

    def _decode(self, code):
        # decode with the codebook
        return np.array([self.codebook_[k] for k in code])

    def encode(self, im, shape=None):
        X, shape = image2data(im, small_size=self.small_size, shape=shape)
        return self.model.predict(X), shape

    def decode(self, code, shape):
        Y = self._decode(code)
        if self.small_size == (1,1):
            Y = Y.reshape((im.size[1], im.size[0])+(3,))
            return Image.fromarray(Y.astype('uint8'))
        else:
            return data2image(Y, small_size=small_size, shape=shape)

    def transform(self, im):
        code, shape = self.encode(im)
        return self.decode(code, shape)

    def stat(self, im, shape=None):
        Y, shape = self.encode(im, shape=shape)
        import collections
        f = collections.Counter(Y)
        N = len(Y)
        p = {k: (v / N) for k, v in f.items()}
        return f, p


def image2data(im, small_size=(1,1), shape=None):
    """Transform an image to 2d array (number of pixels/wh X 3wh)
    where small_size = (w, h)
    
    Arguments:
        im {Image} -- an image
    
    Keyword Arguments:
        small_size {tuple} -- a rectangle domain of an image (default: (1, 1))
        shape {tuple} -- size of small domain (default: {None})
                        calculated based on the size of the image by default
    
    Returns:
        tuple of a block matrix and shape=(r, c)
    """
    w, h = small_size
    if shape is None:
        W, H = im.size
        c, cc = divmod(W, w)
        r, rr = divmod(H, h)
        shape = r, c
        im = im.resize((w*c, h*r))
    else:
        r, c = shape
        im = im.resize((0, 0, w*c, h*r))
    I = np.asarray(im, dtype=np.float64)
    shape = (r, c)
    block = []
    for i in range(r):
        for j in range(c):
            J = I[i*h:(i+1)*h, j*w:(j+1)*w]
            block.append(J.ravel())
    return np.array(block), shape

def data2image(X, small_size, shape):
    # inverse of image2data
    r, c = shape
    w, h = small_size
    X = X.reshape((r, c, w*h*3))
    X = np.block([[[X[i, j].reshape((h, w, 3))] for j in range(c)] for i in range(r)])
    return Image.fromarray(X.astype('uint8', copy=False))
    

def simplevq(im, model=None):
    # vector quantization with a clustering model
    if model is None:
        model = KMeans(n_clusters=2)

    X, r, c = image2data(im, small_size=(2,2))

    model.fit(X)
    code = model.predict(X)  # encoding
    if hasattr(model, 'cluster_centers_'):
        codebook = model.cluster_centers_
    elif hasattr(model, 'means_'):
        codebook = model.means_
    Y = np.array([model.codebook[k] for k in code]) 
    Y = Y.reshape((im.size[1], im.size[0])+(3,))
    return Image.fromarray(Y.astype('uint8', copy=False))

def vq(im, small_size=(2,2), model=None):
    # vector quantization with a clustering model
    if model is None:
        model = KMeans(n_clusters=2)

    X, shape = image2data(im, small_size=small_size)

    model.fit(X)
    code = model.predict(X)  # encoding
    if hasattr(model, 'cluster_centers_'):
        codebook = model.cluster_centers_
    elif hasattr(model, 'means_'):
        codebook = model.means_
    Y = np.array([model.codebook[k] for k in code]) 
    
    return data2image(Y, small_size=small_size, shape=shape)


if __name__ == '__main__':
    model = KMeans(n_clusters=2)
    im = Image.open('cloth.jpg')
    t = VQTransformer(model=model)
    t.fit(im)
    im = t.transform(im)
    im.save('cloth.kmeans2.jpg')
    print(t.stat(im))
    print(t.codebook_)