Kmeans python 可视化

就是用PLT,
上一篇中主要写了kmeans图像分割的算法主体
这里对他做可视化
可视化两个函数，
第一个是visualize，遍历一遍不同的K值的图像，每个暂停0.01s，嘿嘿嘿可以做动画

第二个是交互式的inter_vis 写了两个按钮，+K和-K。

然后为了方便命令行使用又加了parser解析

啊，就是这样

嗯就是这样，UI果然好麻烦。以后还是少写，感觉代码风格都变丑了

from scipy.misc import imread,imshow,imsaveimport numpy as npimport matplotlib.pyplot as pltfrom matplotlib.widgets import Buttonfrom functools import partialfrom optparse import OptionParserdef kmeans(img,K,epsilon):    img = img.astype(np.float64)    randpos = partial(np.random.randint,0,min(img.shape[0],img.shape[1]))    cx,cy = [randpos(K) for i in range(2)]    center = img[cx,cy]    img = img.reshape(1, img.shape[0], img.shape[1], -1)    center = center.reshape(K, 1, 1, 3)    # ite = 0    diff = np.inf    pre_center = np.sum(center)    while(diff>epsilon):        dis = (img - center) ** 2        pos_label = np.sum(dis, axis=3).argmin(axis=0)        for i in range(K): center[i] = np.mean(img[0,pos_label == i],axis=0)        diff = np.abs(np.sum(center)-pre_center)        pre_center = np.sum(center)        # ite+=1        # print(ite,diff)    for i in range(K): img[0,pos_label == i] = center[i]    return np.squeeze(img).astype(np.float16)I = 10def inter_vis(img):    def draw_button():        global button_sub, button_plus        point_sub = plt.axes([0.3, 0.03, 0.1, 0.03])        point_plus = plt.axes([0.6, 0.03, 0.1, 0.03])        button_sub = Button(point_sub, "-")        button_sub.on_clicked(on_press_sub)        button_plus = Button(point_plus, "+")        button_plus.on_clicked(on_press_plus)    def on_press_plus(event):        if event.inaxes == None:            print("none")            return        global I        I = I + 1        fig = event.inaxes.figure        ax1 = fig.add_subplot(111)        nimg = kmeans(img, I, 1e-4)        nimg = nimg.astype(np.uint8)        plt.title("K = " + str(I))        ax1.imshow(nimg)        plt.axis("off")        fig.canvas.draw()    def on_press_sub(event):        if event.inaxes == None:            print("none")            return        global I        if(I==2):            print("at least 2 center!")            return        I = I - 1        fig = event.inaxes.figure        ax1 = fig.add_subplot(111)        nimg = kmeans(img, I, 1e-4)        nimg = nimg.astype(np.uint8)        plt.title("K = " + str(I))        ax1.imshow(nimg)        plt.axis("off")        fig.canvas.draw()    plt.ioff()    fig = plt.figure()    draw_button()    ax1 = fig.add_subplot(111)    nimg = kmeans(img, I, 1e-4)    nimg = nimg.astype(np.uint8)    plt.title("K = " + str(I))    ax1.imshow(nimg)    plt.axis("off")    plt.show()def visualize(img,start=2,end=100):    plt.ion()    plt.axis("off")    for i in range(start,end):        nimg = kmeans(img, i, 1e-4)        nimg = nimg.astype(np.uint8)        plt.title("K = " + str(i))        plt.imshow(nimg)        plt.pause(0.01)if __name__ == '__main__':    parser = OptionParser()    parser.add_option("-v", "--visualize", action="store_true",                      dest="visualize",                      default=False,                      help="visualize kmeans automatically from 2 to 100")    parser.add_option("-i", "--interactive", action="store_true",                      dest="interactive",                      default=False,                      help="use interactive mode to change K")    img = np.floor(imread("/home/ryan/Desktop/cat.jpg"))    (options, args) = parser.parse_args()    if options.visualize == True:        visualize(img)    if options.interactive == True:        inter_vis(img)    if (options.visualize == False) and (options.interactive == False):        visualize(img)