画像セグメンテーション

TensorFlow.orgで表示 GoogleColabで実行 GitHubでソースを表示ノートブックをダウンロードする

このチュートリアルでは、変更されたU-Netを使用した画像セグメンテーションのタスクに焦点を当てます。

画像セグメンテーションとは何ですか?

これまで、画像の分類を見てきました。ネットワークのタスクは、入力画像にラベルまたはクラスを割り当てることです。ただし、オブジェクトが画像内のどこにあるか、そのオブジェクトの形状、どのピクセルがどのオブジェクトに属しているかなどを知りたいとします。この場合、画像をセグメント化する必要があります。つまり、画像の各ピクセルは次のようになります。与えられたラベル。したがって、画像セグメンテーションのタスクは、画像のピクセル単位のマスクを出力するようにニューラルネットワークをトレーニングすることです。これは、はるかに低いレベル、つまりピクセルレベルで画像を理解するのに役立ちます。画像セグメンテーションは、いくつか例を挙げると、医用画像、自動運転車、衛星画像に多くの用途があります。

このチュートリアルで使用されるデータセットは、Parkhi etalによって作成されたOxford- IIITPetDatasetです。データセットは、画像、それに対応するラベル、およびピクセル単位のマスクで構成されています。マスクは基本的に各ピクセルのラベルです。各ピクセルには、次の3つのカテゴリのいずれかが与えられます。

  • クラス1:ペットに属するピクセル。
  • クラス2:ペットに隣接するピクセル。
  • クラス3:上記のいずれでもない/周囲のピクセル。
pip install git+https://github.com/tensorflow/examples.git
import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix

import tensorflow_datasets as tfds

from IPython.display import clear_output
import matplotlib.pyplot as plt

Oxford-IIITPetsデータセットをダウンロードする

データセットはすでにTensorFlowデータセットに含まれているため、ダウンロードするだけで済みます。セグメンテーションマスクはバージョン3以降に含まれています。

dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)

次のコードは、画像を反転する単純な拡張を実行します。さらに、画像は[0,1]に正規化されます。最後に、前述のように、セグメンテーションマスクのピクセルには{1、2、3}のいずれかのラベルが付けられます。便宜上、セグメンテーションマスクから1を引いて、ラベルが{0、1、2}になるようにします。

def normalize(input_image, input_mask):
  input_image = tf.cast(input_image, tf.float32) / 255.0
  input_mask -= 1
  return input_image, input_mask
@tf.function
def load_image_train(datapoint):
  input_image = tf.image.resize(datapoint['image'], (128, 128))
  input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))

  if tf.random.uniform(()) > 0.5:
    input_image = tf.image.flip_left_right(input_image)
    input_mask = tf.image.flip_left_right(input_mask)

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask
def load_image_test(datapoint):
  input_image = tf.image.resize(datapoint['image'], (128, 128))
  input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask

データセットにはすでにテストとトレーニングの必要な分割が含まれているので、同じ分割を引き続き使用しましょう。

TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 64
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE
train = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.AUTOTUNE)
test = dataset['test'].map(load_image_test)
train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
train_dataset = train_dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
test_dataset = test.batch(BATCH_SIZE)

画像の例を見てみましょう。これは、データセットの対応するマスクです。

def display(display_list):
  plt.figure(figsize=(15, 15))

  title = ['Input Image', 'True Mask', 'Predicted Mask']

  for i in range(len(display_list)):
    plt.subplot(1, len(display_list), i+1)
    plt.title(title[i])
    plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
    plt.axis('off')
  plt.show()
for image, mask in train.take(1):
  sample_image, sample_mask = image, mask
display([sample_image, sample_mask])

png

モデルを定義する

ここで使用されているモデルは、変更されたU-Netです。 U-Netは、エンコーダー(ダウンサンプラー)とデコーダー(アップサンプラー)で構成されます。堅牢な機能を学習し、トレーニング可能なパラメーターの数を減らすために、事前にトレーニングされたモデルをエンコーダーとして使用できます。したがって、このタスクのエンコーダーは事前トレーニング済みのMobileNetV2モデルであり、その中間出力が使用されます。 デコーダーは、 Pix2pixチュートリアルのTensorFlowの例ですでに実装されているアップサンプルブロックになります。

3つのチャネルを出力する理由は、各ピクセルに3つの可能なラベルがあるためです。これは、各ピクセルが3つのクラスに分類される複数の分類と考えてください。

OUTPUT_CHANNELS = 3

上述したように、エンコーダを調製し、で使用する準備ができているpretrained MobileNetV2モデルなりtf.keras.applications 。エンコーダーは、モデルの中間層からの特定の出力で構成されます。エンコーダはトレーニングプロセス中にトレーニングされないことに注意してください。

base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)

# Use the activations of these layers
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)

down_stack.trainable = False
Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/mobilenet_v2/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_128_no_top.h5
9412608/9406464 [==============================] - 0s 0us/step

デコーダー/アップサンプラーは、TensorFlowの例で実装された一連のアップサンプルブロックです。

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]
def unet_model(output_channels):
  inputs = tf.keras.layers.Input(shape=[128, 128, 3])

  # Downsampling through the model
  skips = down_stack(inputs)
  x = skips[-1]
  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    concat = tf.keras.layers.Concatenate()
    x = concat([x, skip])

  # This is the last layer of the model
  last = tf.keras.layers.Conv2DTranspose(
      output_channels, 3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

モデルをトレーニングする

あとは、モデルをコンパイルしてトレーニングするだけです。ここで使用されている損失はlosses.SparseCategoricalCrossentropy(from_logits=True)です。この損失関数を使用する理由は、マルチクラス予測と同様に、ネットワークが各ピクセルにラベルを割り当てようとしているためです。真のセグメンテーションマスクでは、各ピクセルに{0,1,2}があります。ここのネットワークは3つのチャネルを出力しています。基本的に、各チャネルはクラスとlosses.SparseCategoricalCrossentropy(from_logits=True)を予測することを学習しようとしていますlosses.SparseCategoricalCrossentropy(from_logits=True)は、このようなシナリオで推奨される損失です。ネットワークの出力を使用して、ピクセルに割り当てられたラベルは、最も高い値を持つチャネルです。これは、create_mask関数が実行していることです。

model = unet_model(OUTPUT_CHANNELS)
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

結果のモデルアーキテクチャをざっと見てみましょう。

tf.keras.utils.plot_model(model, show_shapes=True)

png

モデルを試して、トレーニングの前にモデルが何を予測するかを確認してみましょう。

def create_mask(pred_mask):
  pred_mask = tf.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tf.newaxis]
  return pred_mask[0]
def show_predictions(dataset=None, num=1):
  if dataset:
    for image, mask in dataset.take(num):
      pred_mask = model.predict(image)
      display([image[0], mask[0], create_mask(pred_mask)])
  else:
    display([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])
show_predictions()

png

トレーニング中にモデルがどのように改善されるかを観察してみましょう。このタスクを実行するために、コールバック関数を以下に定義します。

class DisplayCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    clear_output(wait=True)
    show_predictions()
    print ('\nSample Prediction after epoch {}\n'.format(epoch+1))
EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_dataset, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_dataset,
                          callbacks=[DisplayCallback()])

png

Sample Prediction after epoch 20
loss = model_history.history['loss']
val_loss = model_history.history['val_loss']

plt.figure()
plt.plot(model_history.epoch, loss, 'r', label='Training loss')
plt.plot(model_history.epoch, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()

png

予測を行います

いくつかの予測をしてみましょう。時間を節約するために、エポックの数は少なく保たれましたが、より正確な結果を得るには、これを高く設定することができます。

show_predictions(test_dataset, 3)

png

png

png

オプション:不均衡なクラスとクラスの重み

セマンティックセグメンテーションデータセットは非常に不均衡である可能性があります。つまり、特定のクラスのピクセルが他のクラスのピクセルよりも画像の内部に多く存在する可能性があります。セグメンテーションの問題はピクセルごとの分類の問題として扱うことができるため、これを説明するために損失関数を重み付けすることで不均衡の問題に対処できます。これは、この問題に対処するためのシンプルでエレガントな方法です。不均衡なクラスのチュートリアルを参照してください

あいまいさ避けるためにModel.fitは3次元以上の入力のclass_weight引数をサポートしていません。

try:
  model_history = model.fit(train_dataset, epochs=EPOCHS,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            class_weight = {0:2.0, 1:2.0, 2:1.0})
  assert False
except Exception as e:
  print(f"{type(e).__name__}: {e}")
ValueError: `class_weight` not supported for 3+ dimensional targets.

したがって、この場合、自分で均等化を実装​​する必要があります。これは、サンプルの重みを使用して行います。 (data, label)ペアに加えて、 Model.fit(data, label, sample_weight)トリプルも受け入れます。

Model.fitは、 sample_weightを損失とメトリックに伝播します。これらはsample_weight引数も受け入れます。サンプルの重量は、削減ステップの前にサンプルの値で乗算されます。例えば:

label = [0,0]
prediction = [[-3., 0], [-3, 0]] 
sample_weight = [1, 10] 

loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True,
                                               reduction=tf.losses.Reduction.NONE)
loss(label, prediction, sample_weight).numpy()
array([ 3.0485873, 30.485874 ], dtype=float32)

したがって、このチュートリアルのサンプルウェイトを作成するには、 (data, label)ペアを取り、 (data, label, sample_weight)トリプルを返す関数が必要です。ここで、 sample_weightは、各ピクセルのクラスの重みを含む1チャンネルの画像です。

考えられる最も簡単な実装は、ラベルをclass_weightリストへのインデックスとして使用することです。

def add_sample_weights(image, label):
  # The weights for each class, with the constraint that:
  #     sum(class_weights) == 1.0
  class_weights = tf.constant([2.0, 2.0, 1.0])
  class_weights = class_weights/tf.reduce_sum(class_weights)

  # Create an image of `sample_weights` by using the label at each pixel as an 
  # index into the `class weights` .
  sample_weights = tf.gather(class_weights, indices=tf.cast(label, tf.int32))

  return image, label, sample_weights

結果のデータセット要素には、それぞれ3つの画像が含まれます。

train_dataset.map(add_sample_weights).element_spec
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py:5049: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.
Instructions for updating:
The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py:5049: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.
Instructions for updating:
The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.
(TensorSpec(shape=(None, 128, 128, 3), dtype=tf.float32, name=None),
 TensorSpec(shape=(None, 128, 128, 1), dtype=tf.float32, name=None),
 TensorSpec(shape=(None, 128, 128, 1), dtype=tf.float32, name=None))

これで、この重み付けされたデータセットでモデルをトレーニングできます。

weighted_model = unet_model(OUTPUT_CHANNELS)
weighted_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'])
weighted_model.fit(
    train_dataset.map(add_sample_weights),
    epochs=1,
    steps_per_epoch=10)
10/10 [==============================] - 3s 41ms/step - loss: 0.2712 - accuracy: 0.6973
<tensorflow.python.keras.callbacks.History at 0x7fedb45d17d0>

次のステップ

画像セグメンテーションとは何か、そしてそれがどのように機能するかを理解したので、このチュートリアルをさまざまな中間層出力、またはさまざまな事前トレーニング済みモデルで試すことができます。 KaggleでホストされているCarvana画像マスキングチャレンジを試してみることもできます。

また、独自のデータで再トレーニングできる別のモデルのTensorflow Object DetectionAPIを確認することもできます。