numpyで配列を格子状に分割する

CPPXのXです。

画像とかを格子状に分割したかったのですが、どうも関数一発でできないみたいなのでまとめておきます。
そういう関数があったら教えてください。

np.splitで軸を複数指定できればよかったのになあという感じです。

どうにも一発でできないみたいなので、縦分割して...次横分割して...みたいなことをします。

では目次です。

環境

• python 3.5.3
• numpy 1.14.5

imports

import numpy as np

前準備

使うデータを用意します。

t = np.array([
    [1, 2, 3, 4, 5, 6],
    [1, 2, 3, 4, 5, 6],
    [9, 8, 7, 6, 5, 4],
    [9, 8, 7, 6, 5, 4],
    [1, 2, 3, 4, 5, 6],
    [1, 2, 3, 4, 5, 6],
    ]) #[..., np.newaxis] # 6x6x1

6x6のデータです。
6x6x3とか、次元を増やしても大丈夫です。

れっつ分割

分割後の格子数が3x3になるように分割してみます。

np.concatenate(np.hsplit(np.stack(np.hsplit(t, 3)), 3))

これで分割できます(shape: 9x2x2)。
6x2x3にしたい場合は、

np.concatenate(np.hsplit(np.stack(np.hsplit(t, 2)), 3))

こうします。
こんな感じに分割されます。

1 2 3 | 4 5 6
1 2 3 | 4 5 6
------------
9 8 7 | 6 5 4
9 8 7 | 6 5 4
------------
1 2 3 | 4 5 6
1 2 3 | 4 5 6

np.concatenate(np.hsplit(np.stack(np.hsplit(t, 横方向の格子の数)), 縦方向の格子の数))

こうですね。

6x6x3のデータを3x3に分割したら、9x2x2x3になります。

詳細

最初に縦分割します。

np.hsplit(t, 3)

戻り値がリストなので、stackで積み上げます。

np.stack(np.hsplit(t, 3))

この状態で、shapeを見ると3x6x2になります。
縦方向に分割されたものが、横向きに重なっているような状態になっているので、もう一度hsplitで縦分割します。

np.hsplit(np.stack(np.hsplit(t, 3)), 3)

リストが返ってくるので、最初と同様積み上げるのですが、最初に切ったものと区別する必要がないのでconcatenateで連結します。

np.concatenate(np.hsplit(np.stack(np.hsplit(t, 3)), 3))

この状態でshapeを見ると、9x2x2になっていると思います。
分割完了です。

おわり

切ったりくっつけたりしているので、重いような気がします。

速度は確認していません。