[ECE592] Project 4 P1_3 Rate vs Distortion

import cv2
import numpy as np
import matplotlib.pyplot as plt
def runKmeans(img, p, C):

    y , x = img.shape
    number_of_patches = int((y * x)/(p ** 2))
    data = []
    i = 0
    for i in range(0, y - 1, p):
        for j in range(0, x - 1, p):
            patch = img[i:i+p, j:j+p]
            data.append(patch)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.1)
    R = (np.log2(C)*number_of_patches)/(y * x)
    compactness, labels, centers = cv2.kmeans(np.asarray(data).astype('float32'),C,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)
    i = 0
    data_point = 0
    for i in range(0, y - 1, p):
        for j in range(0, x - 1, p):
            img[i:i+p, j:j+p] = centers[labels[data_point][0]].reshape(p,p).astype('uint8')
            data_point += 1
    return [R, compactness**0.5, img]

imgC =  cv2.imread('wolves.png', 0)[27:, 460:972]
p = 2
R = []
D = []
for i in range(0, 7):
    C = (2**i)
    print(C)
    RD = runKmeans(imgC.copy(), p, C)
    cv2.imshow('img',imgC - RD[2])
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    error = np.mean(((imgC.astype('float32') - RD[2].astype('float32')).astype('float32') * (imgC.astype('float32') - RD[2].astype('float32')).astype('float32')).flatten())
    R.append(RD[0])
    D.append(error)

print(R)
print(D)
plt.ylabel('D')
plt.xlabel('R')
plt.plot(R, D)
plt.xticks(range(1, 3))
plt.show()