matplotlibで画像セグメンテーションのマスク画像を表示するとき，同じクラス番号のハズなのに色が変わってしまうときの対処法

解決方法：matplotlib.pyplot.imshowの引数にvmin,vmaxを設定する

matplotlibでは，渡された2次元配列の最大値によって，各値に割り当てるカラーを決めているらしい．
そのため，マスク画像に最大値のクラスが含まれていない場合（4クラス分類なのに，マスク画像に3番の値が含まれていとき），matplotlibの最大値が変わってしまい，割り当てられるカラーが変化してしまう．

そのため，plt.imshowする際に，vmin=0（最小値）とvmas（最大値）=クラス数-1と明示的に指定することで，タグ番号が同じなのに色が変わってしまう問題を解消できる．

例：4クラス分類の場合（画像セグメンテーションの機械学習）

私が直面したトラブルの簡単な例として，画像セグメンテーションのマスク画像をmatplotlibを用いて表示することを例に挙げます．

例えば，以下のようなプログラムでマスク画像を表示するとします．

ex1.py

import matplotlib.pyplot as plt

# マスク画像（4×4 Pixelと仮定）
img = [[0,1,2,3],
       [3,2,1,0],
       [0,1,2,3],
       [3,2,1,0]]

plt.imshow(img)

このとき，以下のような画像が出力されます．

しかし，imgの要素の最大値が2であった場合，

ex1-2.py

import matplotlib.pyplot as plt

# マスク画像（4×4 Pixelと仮定）
img = [[0,1,2,2],
       [2,2,1,0],
       [0,1,2,2],
       [2,2,1,0]]

plt.imshow(img)

出力はこのようになります．要素の3の値がないだけで，それ以外の0,1,2,3の色も最初の例の出力と大きく異なってしまいます．

そこで，plt.imshowの引数にvmin，vmaxを加えて以下のようなプログラムに改良します．

ex1-2_improvement.py

import matplotlib.pyplot as plt

# マスク画像（4×4 Pixelと仮定）
img = [[0,1,2,2],
       [2,2,1,0],
       [0,1,2,2],
       [2,2,1,0]]

plt.imshow(img, vmin=0, vmax=3) # 4クラス分類なので，0~3

すると，0~2の値の色は最初の例ののまま出力することができました．

余談

私は，画像セグメンテーションの推論画像を表示するとき，検出対象のマスク領域は合っているのに，クラス番号が違うといトラブルにに出くわし，困っていました．ずっと学習の仕方が悪いのかと思っていたのですが，matplotlibの仕様でそのように見えていただけなんですね（笑）

対処に大変時間を要しました．灯台下暗しです．

参考