k means on images

kmeans_test.py

def KMeansTest(img,clusters):
    """
    Args:
        path2img   : (str) path to cropped player bounding box
        clusters   : (int) how many clusters to use for KMEANS
    Returns:
        rgb_array :  (tuple) Dominant colors in image in RGB format
    """
    org_img = img.copy()
    #print('Org image shape --> ',img.shape)

    #Convert image into a 1D array
    flat_img = np.reshape(img,(-1,3))
    arrayLen = flat_img.shape
        
    #Do the clustering
    kmeans = KMeans(n_clusters = clusters, random_state=0, tol = 1e-4)
    kmeans.fit(flat_img)
    #Define the array with centroids
    dominant_colors = np.array(kmeans.cluster_centers_,dtype='uint')
    #Calculate percentages 
    percentages = (np.unique(kmeans.labels_,return_counts=True)[1])/flat_img.shape[0]
        
    #Combine centroids representing dominant colors and percentages associated with each centroid into an array
    pc = list(zip(percentages,dominant_colors))
    pc = sorted(pc,reverse=True)
    
    i = 0
    rgb_array = []
    for i in range(clusters):
        dummy_array = pc[i][1]
        rgb_array.append(dummy_array)
        i += 1
            
    return rgb_array