画像の物体検知の予測値をカラーパレット(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]]]