画像の物体検知の予測値をカラーパレット(RGB)値に置換する

備忘録として。

画像の物体検知をディープラーニングを利用して行う際、

学習モデルの予測結果(どの物体と予測したかを示すラベル値)を俯瞰して確認するために、

画像データのドットを予測結果の値に対応するカラーパレット(RGB値)に置き換える操作をしばしば行うらしい(当方、画像処理は無知)。

カラーパレットの数が増えてもコードの修正がないように済む実装方法がないか考えてみた。


動作検証のためのサンプルデータ

import numpy as np

# 学習モデルの予測結果を想定したダミーデータ(0~4までのラベル値)
predImage = np.random.randint(0, 5, (256,512))
print(predImage)

# 0~4のラベルデータに対応するRGB値のカラーパレット
colorPalette = [[0,0,0], [1,1,1], [2,2,2], [3,3,3], [4,4,4]]


置換処理


dotColor = np.reshape([colorPalette[x] for x in predImage.flatten().astype('uint8')], (predImage.shape[0],predImage.shape[1],len(colorPalette[0])))
print(dotColor)


実行結果例


学習モデルの出力した予測結果

[[2 2 3 ... 3 2 3]

[1 4 3 ... 1 1 2]
[1 3 4 ... 4 4 3]
...
[3 4 1 ... 1 2 4]
[1 4 3 ... 3 4 4]
[0 1 1 ... 3 3 1]]


対応するカラーパレットのRGB値に置換した結果

[[[2 2 2]

[2 2 2]
[3 3 3]
...
[3 3 3]
[2 2 2]
[3 3 3]]

[[1 1 1]
[4 4 4]
[3 3 3]
...
[1 1 1]
[1 1 1]
[2 2 2]]

[[1 1 1]
[3 3 3]
[4 4 4]
...
[4 4 4]
[4 4 4]
[3 3 3]]

...

[[3 3 3]
[4 4 4]
[1 1 1]
...
[1 1 1]
[2 2 2]
[4 4 4]]

[[1 1 1]
[4 4 4]
[3 3 3]
...
[3 3 3]
[4 4 4]
[4 4 4]]

[[0 0 0]
[1 1 1]
[1 1 1]
...
[3 3 3]
[3 3 3]
[1 1 1]]]