8x8 vq in python

vq.py

import sys, os, pygame, scipy, numpy as np

# 引数チェック
if len(sys.argv) != 6:
    print("Usage: vq.py <input_file> <w> <h> <k> <output_file>")
    sys.exit(1)
inFile, W, H, K, outFile = sys.argv[1:]
W,H,K=[int(v) for v in [W,H,K]]

if not os.path.isfile(inFile):
    print("{} is not a valid file".format(inFile))
    sys.exit(2)

# 1. 画像読込
img = pygame.surfarray.array3d(pygame.image.load(inFile))
h, w = (img.shape[0], img.shape[1]) # サイズ取得

# 2. ベクトル化する
N = W*H*3                    # ベクトルの次元 8x8x3=192
qImg = np.empty(((w * h) // (N//3), N), dtype = np.uint8) # ベクトル化された画像
i = 0
for y in range(0, h, H):
    for x in range(0, w, W):
        qImg[i] = np.array(img[y : y + W, x : x + H]).reshape([N])
        i += 1

# 3. kmeansで量子化
ptns, names = scipy.cluster.vq.kmeans2(np.array(qImg, dtype=np.float64),
                    k = K, iter = 300, thresh = 0.0001, minit = '++')

# 4. パターンテーブルとネームテーブルから画像を作る
img = np.empty((h,w,3), dtype=np.uint8)
i = 0
for y in range(0, h, H):
    for x in range(0, w, W):
        img[y : y + W, x : x + H] = ptns[names[i]].reshape([W, H, 3])
        i += 1

# 5. 画像出力
pygame.image.save(pygame.surfarray.make_surface(img),outFile)