Edited at

Segmentation用のInputに対するTips


Segmentationについて

Deep Oculusの@hmasumotoです。今回はSegmentationにおけるデータの取り扱いについて、紹介したいと思います。

医学におけるAI論文で最も使用されている技術はSegmentationであります。

実際に、画像を解析する際には、DatasetとしてすでにTensorにされているものではなく、jpgなどで保存されている画像をもとにnumpyのArrayを作っていく必要があります。

ここで、実際のSegmentation画像をnumpyのarrayに変換していく方法を紹介したいと思います。

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

sample_png.png

上記、画像は完全な黒色([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数の形に変換する必要があります。

sample_png.png

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

ここで、今回紹介しますコードの条件として、必ず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などを開発していきたいと思っています。

非常に実践的な内容を、紹介していきたいので、今後もよろしくお願いいたします。