compression_kmeans.py

# -*- coding: utf-8 -*-
"""
Created on Wed Apr 11 11:01:05 2018

@author: 3775548
"""

import  matplotlib.pyplot  as  plt
import numpy as np
import random
from matplotlib.colors import hsv_to_rgb
from matplotlib.colors import rgb_to_hsv

def get_prototype_barycentre(C):
    mu = np.zeros(3)
    count = 0
    
    for x in C:
        count += 1
        mu += x
        
    return mu / count if count != 0 else mu
    
def get_distance(mu, x):
    v = mu - x
    return np.sqrt(v[0]*v[0] + v[1] * v[1] + v[2] * v[2])

def get_compression_dict(pixels, n_protos):
    done = False
    
    size = pixels.size / 3
    print(size)
    orig_indices = [i for i in range(size)]
    print("mooh")
    protos = []
    for i in range(n_protos):
        protos.append([])    
    print("boo")
    for i in range(n_protos):
        for j in range(size / n_protos):
            pick = random.choice(orig_indices)
            orig_indices.remove(pick)
            protos[i].append(pixels[pick])
            
    for p in orig_indices:
        protos[0].append(pixels[p])
            
    barycentre_list = [0 for _ in range(n_protos)]
    print("done")
    done = False
    for m in range(10):
        print("hey")   
        
        done = True
        i = 0
        
        for proto in protos:
            new_barycentre = get_prototype_barycentre(proto)
            print(new_barycentre)
            print(barycentre_list[i])
            if not all(new_barycentre == barycentre_list[i]):
                done = False
            barycentre_list[i] = new_barycentre
            i += 1
            
        print(barycentre_list)
            
        if done:
            break
        
        protos = []
        for i in range(n_protos):
            protos.append([])
            
        for c in pixels:               
            min_dist_proto = 0
            min_dist = -1
                
            for i in range(n_protos):
                dist = get_distance(barycentre_list[i], c)
                
                if min_dist > dist or min_dist == -1:
                    min_dist = dist
                    min_dist_proto = i
                        
            protos[min_dist_proto].append(c)
            
            for p_id in range(len(protos)):
                if len(protos[p_id]) == 0:
                    min_dist_pixel_id = 0
                    min_dist = -1
                    for i in range(pixels.size / 3):

                        dist = get_distance(barycentre_list[p_id], pixels[i])                        
                        
                        if min_dist > dist or min_dist == -1:
                            min_dist = dist
                            min_dist_pixel_id = i
                            
                            
                    protos[p_id].append(pixels[min_dist_pixel_id])
            
    res = dict()
    for i in range(len(protos)):
        for e in protos[i]:
            res[tuple(e)] = barycentre_list[i]
    return res
        

im=plt.imread ("Lenna.png" ) [ : , : , : 3 ] #on  garde  que  l e s  3  premieres  composantes ,  la  transparence  est  i n u t i l e
im_h, im_l ,_=im.shape
im = rgb_to_hsv(im)
pixels = im.reshape ((im_h * im_l , 3 ) ) #transformation  en  matrice n * 3 , n nombre de  p i x e l s
imnew=pixels.reshape ((im_h, im_l , 3 ) ) #transformation  inverse
plt.imshow (im) #a f f i c h e r  l ’ image
p_dict = get_compression_dict(pixels, 60)

for i in range(im_h):
    for j in range(im_l):
        im[i][j] = p_dict[tuple(im[i][j])]
        
plt.imshow (im) #a f f i c h e r  l ’ image