Numpyでインデックスカラー画像（VOC2012のマスク）→RGB画像への変換をする方法

---

Semantic Segmentationのマスク画像には「インデックスカラー」というRGBとは異なったフォーマットを使っていることが多いです。この形式はPILで扱うことができ、RGBに変換できます。しかし、モデルの予測を表示したいときはダイレクトにNumpyで計算できると便利です。VOC2012のマスク画像を例に見ていきます。

インデックスカラーとは

通常のRGBによるカラー画像は、(x, y)の点に対してRGBのピクセル値を格納しています。インデックスカラーは(x, y)座標を使っていることは変わらないのですが、RGBの具体的な値ではなくカラーパレットのIDを格納しています。

例えばカラーパレットが、

• 0 : RGB=(0, 0, 0) 黒
• 1 : RGB=(255, 0, 0) 赤
• 2 : RGB=(0, 255, 0) 緑
• 3 : RGB=(0, 0, 255) 青

という定義で、2×2の画像の左上から青→緑→赤→黒という画像にしたい場合、インデックスカラー画像では次のようになります。

[[3, 2],
 [1, 0]]

もしこれがRGB画像の場合は次のようになります。

[[[0, 0, 255], [0, 255, 0]],
 [[255, 0, 0], [0, 0, 0]]]

インデックスカラーのほうがぱっと見てわかりやすいですね。Numpyの配列にすると、RGB画像が(Height, Width, 3)という3ランクのテンソルになりますが、インデックスカラー画像は(Height, Width)という1ランクのテンソルになります。

なぜこのインデックスカラーがSemantic Segmentationで使われるのかというと、パレットのインデックスがそのままクラスインデックスに対応できるからです。つまり普段見るような、

のような画像は、インデックスカラーをあたかもRGB画像のように表示しているだけということになります。

コード

VOC2012の最初の画像を使いました。torchvisionから読み込ませます。

import torchvision
from PIL import Image
import numpy as np

# PILのimg(mode=P).convert("RGB")と同じ処理
def convert_to_rgb(index_colored_numpy, palette, n_colors=None):
    assert index_colored_numpy.dtype == np.uint8 and palette.dtype == np.uint8
    assert index_colored_numpy.ndim == 2 and palette.ndim == 2
    assert palette.shape[1] == 3
    if n_colors is None:
        n_colors = palette.shape[0]
    reduced = index_colored_numpy.copy()
    reduced[index_colored_numpy > n_colors] = 0 # 不要なクラスを0とする
    expanded_img = np.eye(n_colors, dtype=np.int32)[reduced]  # [H, W, n_colors] int32
    use_pallete = palette[:n_colors].astype(np.int32)  # [n_colors, 3] int32
    return np.dot(expanded_img, use_pallete).astype(np.uint8)

def main():
    # VOC2012はダウンロード済みとする（未ダウンロードの場合は download=True）
    trainset = torchvision.datasets.VOCSegmentation(root="./data", download=False, image_set="train")
    path = next(iter(trainset.masks))

    with Image.open(path) as img:
        palette = np.array(img.getpalette(), dtype=np.uint8).reshape(-1, 3) # パレットの取得
        p_array = np.asarray(img)  # Numpy配列に変換
        print(p_array.shape)
        converted_rgb = np.asarray(img.convert("RGB")) # PILでコンバート
        rgb_array = convert_to_rgb(p_array, palette)
        print(rgb_array.shape)
        print(np.all(converted_rgb == rgb_array))  # PILでのコンバートと結果が等しいか確認用
    with Image.fromarray(rgb_array) as img:
        img.save("out.png")

if __name__ == "__main__":
    main()

PILでは「img.convert(“RGB”)」とすることで、RGBに変換できます（逆変換はバグなのかうまくいかないので要注意、どうしても逆変換をしたい場合はquantizeを使ったほうが無難）

やっていることは、PILのimg.convert(“RGB”)と、Numpyだけで計算したコードの出力が一緒かを確かめています。

Numpyだけの計算はconvert_to_rgbの部分です。言っちゃえばこれは行列の内積を取るだけなので、np.dotで終わります。出力は次のようになります。

(281, 500)
(281, 500, 3)
True

上から、インデックスカラーの画像をNumpy化したときのshape、RGB画像をNumpy化したときのshape、PILのimg.convert(“RGB”)とconvert_to_rgbの出力が全ピクセルで同じかを比較した結果です。

不要なクラスを省いてみる

ちなみにこのconvert_to_rgbは使うパレット数を制限できるので、VOCの仕様に合わせて使うパレットを21（背景＋20クラス）に制限してみます。こうすればいいだけです。

convert_to_rgb(p_array, palette, n_colors=21)

境界線（undifinedの領域）が消えました。

Numpy配列からインデックスカラーの画像を作る

逆の方法として、Numpy配列からインデックスカラーの画像を作ってみましょう。

def indexed_image_from_scratch():
    indices = (np.arange(350) // 50).astype(np.uint8).reshape(1, -1)  # [1, 350] px
    indices = np.broadcast_to(indices, (50, 350))  # [50, 350] px

    # 虹色のカラーパレット
    color_palette = [
        255, 0, 0,
        255, 165, 0,
        255, 255, 0,
        0, 128, 0,
        0, 255, 255,
        0, 0, 255,
        128, 0, 128
    ]
    with Image.fromarray(indices, mode="P") as img:
        img.putpalette(color_palette)  # パレットの設定
        img.save("rainbow.png")

虹色を50px置きに描画する関数です。結果はこのようになります。

うまくできました。Image.fromarrayでグレー画像のようにロードしたあと（ただし各値はインデックス）、パレットをputpaletteで設定すればOKです。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---