はじめに

NumPyを用いて3次元のデータ配列を操作する際、データとして(フレーム番号、画像の縦、画像の横)の形状を持つ場合があります。このようなデータを効率的に操作し、特定のフレームの特定の画素にアクセスする方法は、データ解析において必要だと思います。本記事では、可読性を損なわずに簡潔にコードを書く方法を紹介します。

具体的には、画素毎に任意のフレーム番号が与えられた場合に、forループなどを使わずにデータを迅速に抽出する方法をサンプルコードを交えて説明します。

方法

サンプルでは簡単のため、最大値のインデックスの配列を使って、対応したインデックスの画素値を取得する方法について紹介します。

本記事で使用したコードは、Google Colabのこちらからも閲覧できます。

方法1

はじめにシンプルな方法について紹介します。

# 方法1
import numpy as np


# 再現性のためseedを固定
np.random.seed(2023)

# 配列を作成
array = np.random.randint(0, 10, (3, 2, 4))

# 最大値のインデックス配列を取得
index_array = np.argmax(array, axis=0)

# グリッド配列の生成
_, h, w = array.shape
W, H = np.meshgrid(np.arange(w), np.arange(h))

# インデックス配列に対応したデータの抽出
get_array = array[index_array, H, W]
# get_array = np.max(array, axis=0)  # 今回はこの書き方と同じ結果になる

print("元の配列:\n", array)
print("インデックス配列:\n", index_array)
print("インデックス配列に対応したデータ配列:\n", get_array)

出力結果

元の配列:
 [[[7 9 6 7]
  [1 3 4 4]]

 [[6 5 0 6]
  [1 5 7 5]]

 [[8 2 3 1]
  [0 7 1 0]]]
インデックス配列:
 [[2 0 0 0]
 [0 2 1 1]]
インデックス配列に対応したデータ配列:
 [[8 9 6 7]
 [1 7 7 5]]

コードの説明

最初に、NumPyライブラリをインポートします。
次に、再現性のために乱数のシード（seed）を固定します。
配列arrayを作成します。この配列は(3, 2, 4)の形状を持ち、各要素には0から9までのランダムな整数が含まれています。
np.argmax関数を使用して、arrayの各画素において最大値を持つインデックスを取得します。axis=0を指定することで、各画素における最大値のインデックスを縦と横の軸に対して求めます。結果はindex_arrayに格納されます。このインデックスを使用して、特定のフレームに対応する画素値を取得することができます。
画素の座標を格納するためのグリッド配列HとWを生成します。
※ NumPyのバージョンが1.17以上であればindices()が利用可能で、meshgrid()よりもこの場面に適していると思います。具体的には、W, H = np.meshgrid(np.arange(w), np.arange(h))の箇所を、H, W = np.indices((h, w))として使います。
最後に、arrayから特定のフレームに対応する画素値を取得するために、array[index_array, H, W]を使用します。これにより、各画素の座標に対応するデータが抽出されます。

このコードは、簡単のため最大値を使っており、実際はget_array = np.max(array, axis=0)と同じ操作を行っていることになり、このコードだけではあまり実用的に思えないかもしれません。

しかし、この記事のポイントは、index_arrayを任意で使えることです。活用例として、時系列画像が2つあったとして、片方の最大値のインデックスに対応したもう片方の配列からデータを抽出する時などに役に立つと思います。

方法2

np.take_along_axis関数を使っても同様の処理をより簡潔に書くことができます。

np.take_along_axisの使い方は、本記事の目的以外にも以下の記事など参考になります。

# 方法2
import numpy as np


# 再現性のためseedを固定
np.random.seed(2023)

# 配列を作成
array = np.random.randint(0, 10, (3, 2, 4))

# 最大値のインデックス配列を取得
index_array = np.argmax(array, axis=0, keepdims=True)


# インデックス配列に対応したデータの抽出
get_array = np.take_along_axis(array, index_array, axis=0)[0]
# get_array = np.max(array, axis=0)  # 今回はこの書き方と同じ結果になる

print("元の配列:\n", array)
print("インデックス配列:\n", index_array)
print("インデックス配列に対応したデータ配列:\n", get_array)

出力結果

方法1と同じ

コードの説明

np.argmax 関数を使用して、arrayの各画素において最大値を持つインデックスを取得します。axis=0およびkeepdims=Trueを指定することで、各画素における最大値のインデックスを縦と横の軸に対して求め、結果はindex_arrayに格納されます。
この段階で、index_arrayは次元を保持するため、そのままデータ抽出に使用できません。そのため、np.take_along_axis関数を使用して、arrayから特定のフレームに対応する画素値を抽出します。axis=0により、index_arrayの値に対応するデータを選択します。次元を1つ落とすため末尾に[0]を加え、get_arrayに格納します。

この方法はnp.take_along_axis関数に慣れる必要があり、私もあまり詳しくないためChatGPTに聞きながらコードを作っていますが、可読性は良いと思います。

まとめ

この記事では、NumPyを活用して3次元の時系列画像データから、任意のインデックス配列を使用してデータを効率的に抽出する方法について解説しました。

方法1では、基本的なインデックスを活用してデータを抽出する方法を紹介しました。このアプローチはシンプルかつ理解しやすい一方、np.take_along_axis関数を使用する方法と比べて、メッシュ配列の用意など一部手間がかかります。

方法2では、np.take_along_axis関数を活用し、最大値のインデックスを保持しながらデータを抽出する方法を説明しました。この方法はコードが簡潔であり、本記事だけでなく、さまざまな応用シナリオで利用できる汎用性が高いです。

最後に、他にも有用な実装方法が存在するかもしれません。読者の方で新たな方法をご存知の場合は、ぜひ共有いただければ幸いです。

この記事が、NumPyを使用したデータ処理に新たなアプローチを提供し、読者のプロジェクトや研究に役立つことを願っています。