Kmeans python 可视化

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

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

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

啊，就是这样

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

from scipy.misc import imread,imshow,imsave
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button
from functools import partial
from optparse import OptionParser

def 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 = 10
def 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)