Pythonで簡単な画像をセグメンテーションした後に領域の周長を求める話

※全文無料です。

1. 画像データセットの用意

とりあえずセマンティックセグメンテーションを行うためのデータセットを用意しましょう。64*64ピクセルの黒背景の上に、['▲','●','■']の中からランダムに2つ配置して画像化します(返すのはndarrayですが)。片方にはノイズを乗せておきましょう。メモリが不安なら画像として出力処理を入れてください。

from PIL import Image, ImageDraw, ImageFont
import numpy as np

symbols = np.array(['▲','●','■'])

def im_gen():
   font = ImageFont.truetype('times.ttf', np.random.randint(12, 20))
   im = Image.new("L",(64,64),0)
   draw = ImageDraw.Draw(im)
   draw.text(np.random.randint(0, 20, 2), ''.join(np.random.choice(symbols)), font=font, fill=255)
   draw.text(np.random.randint(25, 50, 2), ''.join(np.random.choice(symbols)), font=font, fill=255)
   return np.array(im)

def noise(x):
   y = np.where(x < 100, x + 100, x) - np.random.randint(0, 100, 64*64).reshape(64,64)
   return y

label_data = []
image_data = []

for _ in range(10000):
   im = im_gen()
   label_data.append(im)
   image_data.append(noise(im))

label_data = np.array(label_data)
image_data = np.array(image_data)

さて、これを実行すると次のような画像のペアが1万組出力されます。Image.fromarray()で1枚ずつ見ることができます。

Image.fromarray(np.uint8(label_data[0]))
Image.fromarray(np.uint8(image_data[0]))

これが確認できればデータセットの作成は終了です。

2. セマンティックセグメンテーション

TensorFlowでノイズ付きの画像からノイズ無しの画像を取り出すニューラルネットワークを構成します。その前にまずはデータセットを少し加工しましょう。

import tensorflow as tf

images = tf.reshape(image_data/255.0, [-1, 64, 64, 1])
labels = tf.reshape(label_data/255.0, [-1, 64, 64, 1])

これでさきほど出力したndarrayが全てテンソルになりました。ついでにピクセルごとの輝度も0~1の範囲に収まるように255で割って正規化しています。
次にモデルを組んでいきましょう。

im_input = tf.keras.layers.Input(shape=(None, None, 1))

conv1 = tf.keras.layers.Conv2D(16, (5, 5), padding='same')(im_input)
conv1 = tf.keras.layers.BatchNormalization()(conv1)
conv1 = tf.keras.layers.Activation('relu')(conv1)

conv2 = tf.keras.layers.Conv2D(32, (5, 5), padding='same')(conv1)
conv2 = tf.keras.layers.BatchNormalization()(conv2)
conv2 = tf.keras.layers.Activation('relu')(conv2)

train_out = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', padding='same')(conv2)

model = tf.keras.models.Model(inputs = im_input, outputs = train_out)

model.summary()

インプット層の下に畳み込み層→Batch Normalization→活性化関数ReLUの順で2回並べて、2値分類なのでアウトプット層は活性化関数Sigmoidを使っています。組んだモデルをコンパイルします。

adam = tf.keras.optimizers.Adam(lr=0.0001, epsilon=1e-06)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=[tf.keras.metrics.binary_accuracy])

2値分類なのでロス関数はbinary crossentropyです。オプティマイザーはAdamで問題ないでしょう。

import datetime

model.fit(x=images, y=labels, batch_size=16, epochs=3, verbose=1, validation_split=0.1, shuffle=True, steps_per_epoch=9000//16, validation_steps=1000//16)
model.save('saved_model/' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))

#Epoch 1/3
#562/562 [==============================] - 40s 71ms/step - loss: 0.0059 - binary_accuracy: 0.9894 - val_loss: 0.0054 - val_binary_accuracy: 0.9894
#Epoch 2/3
#562/562 [==============================] - 39s 70ms/step - loss: 0.0054 - binary_accuracy: 0.9895 - val_loss: 0.0053 - val_binary_accuracy: 0.9894
#Epoch 3/3
#562/562 [==============================] - 40s 71ms/step - loss: 0.0053 - binary_accuracy: 0.9894 - val_loss: 0.0053 - val_binary_accuracy: 0.9894
#INFO:tensorflow:Assets written to: saved_model/20210309-132235\assets

実際に学習させた後にモデルをセーブします。データが単純なのでエポック数は3もあれば十分です。もし、ここでプログラムが動かない場合には、

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

とすればGPUが無効になってCPUのみを使うようになるので、まずはこれを試してみましょう。これでダメならメモリが足りていない可能性が高いです。
では、モデルを使ってセグメンテーションしてみましょう。

pred_data = []

for _ in range(20):
   im = im_gen()
   pred_data.append(noise(im))

pred_data = np.array(pred_data)

predictions = model.predict(pred_data/255.0)

Image.fromarray(np.uint8(predictions.reshape(-1,64,64)[0]*255))
Image.fromarray(np.uint8(pred_data[0]))

とすれば、ノイズ画像とセグメンテーション後の画像のペアが得られます。

今回は単純なデータで行いましたが、畳み込みニューラルネットワークを用いた全てのセマンティックセグメンテーションの基本的な流れはこれと同じです。どんどん使っていきましょう。

3. 領域の周長を求める

まず、今後の説明を楽にするために、セグメンテーション後の画像を白黒の2値化しましょう。

def w_or_b(x):
   y = np.where(x < 127, 0, 255)
   return y

pred_image = np.uint8(w_or_b(predictions.reshape(-1,64,64)[0]*255)/255)

さて、今回求めたいものは2つの図形の周長を足した合計です。そのために、まずは画像の元となる配列に前処理を行います。

pred_image = np.repeat(np.repeat(pred_image, 2, axis=1), 2, axis=0)
pred_image = np.block([np.zeros((128, 1), dtype='uint8'), pred_image])
pred_image = np.block([pred_image, np.zeros((128, 1), dtype='int16')])
pred_image = np.insert(pred_image, 0, np.zeros((1, 128+2), dtype='int16'), axis=0)
pred_image = np.insert(pred_image, 128+1, np.zeros((1, 128+2), dtype='int16'), axis=0)

画像を縦横2倍に引き延ばした後、画像の外側に黒色のピクセルを1ピクセルずつ追加しています。こうすることで、境界が重なっているときや、画像のふちに図形が重なっているときでも領域が求まります。
続いて輪郭を定義します。

for i in range(1, 128+2):
   for j in range(1, 128+2):
       count=0
       if pred_image[j][i]==0 and pred_image[j+1][i]==1:
           count+=1
       if pred_image[j][i]==0 and pred_image[j][i+1]==1:
           count+=1
       if pred_image[j][i]==0 and pred_image[j-1][i]==1:
           count+=1
       if pred_image[j][i]==0 and pred_image[j][i-1]==1:
           count+=1
       if count>=1:
           pred_image[j][i]=-1
       if count==2:
           pred_image[j][i]=-2

わかりやすいようにfor文の2重ループで書きましたが、実際に使う場合には通常通り速い表記を使ってください。このプログラムでは、背景ピクセルの隣に1ピクセル白色があるときに配列の要素を-1とし、隣に2ピクセル白色があるときに配列の要素を-2とします。
このようにすれば、pred_imageの0未満の要素を合計して絶対値を取ったものの半分が、2つの図形の周の長さの合計になります。

r = abs(sum(pred_image[pred_image<0]))//2
# r = 82