Segmentationについて

Deep Oculusの@hmasumotoです。今回はSegmentationにおけるデータの取り扱いについて、紹介したいと思います。
医学におけるAI論文で最も使用されている技術はSegmentationであります。
実際に、画像を解析する際には、DatasetとしてすでにTensorにされているものではなく、jpgなどで保存されている画像をもとにnumpyのArrayを作っていく必要があります。
ここで、実際のSegmentation画像をnumpyのarrayに変換していく方法を紹介したいと思います。

ここで、この白黒画像を題材として、紹介していきたいと思います。

上記、画像は完全な黒色（[0,0,0]）と完全な白色（[255,255,255]）のみからなっています。つまり、中間色はないということです。

1. まず、そもそもDatasetを作るときの注意

JPGは非可逆圧縮の拡張子です。
先に結論を言うのであれば、JPEGを使うな、PNGなどの可逆圧縮拡張子で保存すべきということです。
理由をこれから説明します。

def count_middle(path):
    img_png = cv2.imread(path)
    print(np.sum((img_png > 0 )& (img_png < 255)))

上記コードは、中間色（0と255以外の値）のpixel数をカウントするコードです。
これを用いて上記画像のpngファイルを用いて、中間色をカウントすると、

count_middle("XXX.png")

無事に0が出力されました。

しかし、一旦、pngの画像を読み込み、jpgファイルに保存したあと、再度中間色の数を

img = cv2.imread("XXX.png")
cv2.imwrite("XXX.jpg", img)
img2 = cv2.imread("XXX.jpg")
count_middle("XXX.jpg")

を行うと、3384という数字が出力されてしまいました。
ということで、
特定の色を用いてクラスを記述していく、
つまり、
[255,255,255]であるpixelをクラス1とするみたいな方法
を使う場合には、
__jpgを用いると、クラスがきちんと読み込めなくなってしまう__ことが分かります。

2. BGR(RGB)↔Classの変換

まず、画像を読み込んだ際のNumpyのArrayはheight * width * 3ch(BGR)の形になっています。
しかし、これでは、SegmentationのNeural NetworkにInputするArrayの形はheight * width * class数の形に変換する必要があります。

ここで今回は、画像のうち、黒い部分をclass 0 (背景)、白い部分をSegmentationしたいAreaであるclass 1 (物体)という２クラス分類を行います。

ここで、今回紹介しますコードの条件として、必ずclass 0 (背景)は必ず([0,0,0]の真っ黒であることを前提とします。)

では、さっそくコードを紹介します。

color_dic = {1:[255,255,255]}

もし、class 0 ([0,0,0])、class 1(赤)([0,0,255]), class 2(青)([255,0,0])とするのであれば、

color_dic = {1:[0,0,255], 2:[255,0,0]}

とします。（OpenCVを使う前提なので、BGR形式で記載しています。）

normal_denormal_labeling.py

import cv2
import os
import numpy as np
class Normalize():
    def __init__(self,color_dic):
        self.color_dic = color_dic
    def color_pick(self,img,color):
        ch1 = np.zeros((img.shape[0],img.shape[1]))
        ch1[(img[:,:,0] == color[0])&(img[:,:,1] == color[1])&(img[:,:,2] == color[2])] = 1
        return ch1
    def normalize(self,img):
    # labelingの箱をつくる
        outputs = np.zeros((img.shape[0],img.shape[1],len(self.color_dic) + 1))
        # label1 - の中に、条件を満たすピクセルのみ1に変換していく
        for i in range(1,len(self.color_dic)+1):
            ch1 = self.color_pick(img,self.color_dic[i])
            outputs[:,:,i] = ch1
        # 1番上の層（つまり、背景以外が0のもの）のみ抽出
        ch1sum = np.sum(outputs[:,:,1:] , axis = 2)
        # 一番上の層に代入する板を用意
        ch0 = np.zeros((img.shape[0],img.shape[1]))
        # 背景以外が0のもののみを1に変換する
        ch0[ch1sum == 0] = 1
        # 1番上の層の中身を代入する。
        outputs[:,:,0] = ch0
        return outputs
class DeNormalize():
    def __init__(self,color_dic):
        self.color_dic = color_dic
    # tagは、予測されたheight * width * classのArray
    def denormalize(self, tag):
        tag_arr = np.argmax(tag,axis = 2)
        img = np.zeros((tag.shape[0],tag.shape[1],3),dtype = np.uint8)
        for i in range(len(self.color_dic)+1):
            if i == 0:
                for idx in range(3):
                    img[:,:,idx][tag_arr == i] = 0
            else:
                for idx in range(3):
                    img[:,:,idx][tag_arr == i] = self.color_dic[i][idx]
        return img

このNormalizeクラスがBGRのArray→ClassのArrayに変換するクラスです。

norm = Normalize(color_dic)
norm.normalize(img)

みたいな使い方をします。

このDenormalizeクラスがClassのArray→BGRのArrayに変換するクラスです。

denorm = Denormalize(color_dic)
denorm.denormalize(tag)

みたいな使い方をします。

最後に

初投稿でしたが、いかがでしたか？
医療機関でありながら、本格的な開発を行えるプログラマーチームとして、今後も医療AIなどを開発していきたいと思っています。
非常に実践的な内容を、紹介していきたいので、今後もよろしくお願いいたします。