チャンネル数の変更 深層学習

Question

深層学習でU-Netを用いてセマンティックセグメンテーションをしようとしています。
物体クラスと背景クラスの２クラスなのでU-Netの最終層のチャネル数は2です。

出力データのshapeが(2, 256, 256)なのですが、このセグメンテーションマスクを画像として保存するときにおそらく2チャンネルが原因でエラーになります。
保存にはcv2.imwriteを用いています。

shapeが(3, 256, 256)なら保存できるのでチャネル数を2から3にしたいのですが、どうしたらよいでしょうか？

調べてもグレースケールからRGB変換しか見つかりません。

よろしくお願いいたします。

Comments

[jbpb0]

3チャンネルの画像で保存したいというのは、たとえば背景は緑で物体は赤、みたいな画像にしたいのでしょうか？

そうするよりも、shapeが(256, 256)のグレースケールにして、各画素に0または1を代入するようにした方が、クラスを3よりも多くしたいときにそのまま使えるので汎用性があると思うのですが

Answer 1

得られた出力を ndarray に変換したものを pred とします。

pred = np.argmax(pred, 0) # (2, 256, 256) -> (256, 256) 要素は 0 or 1
output = np.zeros((pred.shape[0], pred.shape[1], 3)) # (256, 256, 3) の画像（黒）を作成
output[pred==1] = [255, 255, 255] # 0は黒、1は白とする

# 画像の出力
cv2.imwrite('output.png', output)

コードの確認はしていないのでどこか間違いがあるかもしれませんが大まかな流れとして、
・得られ出力を argmax とって 0,1 にする
・同じサイズで3軸目を3にした空(0)の配列を用意
・先ほど argmax とったもので1となっている個所を白[255,255,255]に置き換える
・画像出力

補足として、これが3クラス分類になったら、argmax 後の値が 0,1,2 となるので
output[pred==0] = 0と推測したピクセルのBGR
output[pred==1] = 1と推測したピクセルのBGR
output[pred==2] = 2と推測したピクセルのBGR
とすれば希望の色に変更できます。
また、出力方法に OpenCV を用いるのであれば BGR、Pillow を用いるのであれば RGB なので注意です。