fastaiを用いたセマンティックセグメンテーション

概要

今回紹介するnotebookはAerial_Drone_semantic_image_segmentation | Fastaiです。 このnotebookの内容はおおまかに以下の通りです。

  • セマンティックセグメンテーション(semantic segmentation)とは?
  • fastaiを使ってセマンティックセグメンテーションをやってみる

セマンティックセグメンテーションとは?

セマンティックセグメンテーションとは画像処理技術の一つで、画像内の物体をピクセルレベルで分類していくものです。 イメージは以下のようなものです。 semantic-segmentation

引用:A Beginner’s guide to Deep Learning based Semantic Segmentation using Keras

上記の例では、ベッドに属するピクセルはクラス「ベッド」に分類され、壁に対応するピクセルには「壁」という風にラベルが付けられます。 このように、サイズW x H x 3の画像に対して、すべてのピクセルに予測クラスラベルを含むW x H行列を生成することが目標です。 以下の図がそのイメージです。 semantic-segmentation

引用:A Beginner’s guide to Deep Learning based Semantic Segmentation using Keras

やってみる

それでは早速コードを辿っていきましょう。 今回使っているデータセットはAerial Semantic Segmentation Drone DatasetDrone_images_mask_resized の2つを利用しています。

ライブラリインストール

from fastai import *
from fastai.vision import *
from fastai.callbacks.hooks import *
from fastai.callbacks import SaveModelCallback
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.image as immg
import gc
import numpy as np
import random
from PIL import Image
import warnings
warnings.filterwarnings("ignore")
sns.set_style('darkgrid')

データセットにパスを通す

path = Path('../input/drone-images-mask-resized/drone_data_small')
path.ls()

semantic-segmentation

fnames = get_files(path/'train_small')
fnames_mask = get_files(path/'label_small')

データ前処理

from tqdm.notebook import tqdm,tnrange

path_im = path/'train_small'
path_lb = path/'label_small'

get_y_fns = lambda x: path_lb/f'{x.stem}.png' #画像のマスクをする関数

def get_classes(fnames):
    class_codes=[]
    for i in tqdm(range(400)):
        class_codes += list(np.unique(np.asarray(Image.open(get_y_fns(fnames[i])))))
    return np.array(list(set(class_codes)))

codes = get_classes(fnames) #クラスのリストを取得
codes = np.array(codes)
print(codes)

semantic-segmentation

画像をマスキングする関数

def drone_mask(f):  # f = file_name
  img_a = immg.imread(f)
  img_a_mask = immg.imread(get_y_fns(f))
  plt.figure(1,figsize=(20,8))
  plt.subplot(121)
  plt.imshow(img_a);plt.title('Raw Drone footage ');plt.axis('off')
  plt.subplot(122)
  plt.imshow(img_a,alpha=0.8);
  plt.imshow(img_a_mask,alpha=0.8);plt.title('Drone with  mask');plt.axis('off')
  plt.show()

サンプルとして画像をマスキングしてみる

img_num = random.randint(10,200)
drone_mask(fnames[img_num])

semantic-segmentation

学習のためのデータセット作成

src = np.array([400,600])

data = (SegmentationItemList.from_folder(path=path_im)
        .split_by_rand_pct(0.2)
        .label_from_func(get_y_fns, classes=codes)
        .transform(get_transforms(), size=src//2, tfm_y=True)
        .databunch(bs=4)
        .normalize(imagenet_stats))

data.show_batch(rows=1,figsize=(10,5))

semantic-segmentation


モデル

name2id = {v:k for k,v in enumerate(codes)}
void_code = -1

def drone_accuracy_mask(input, target):
    target = target.squeeze(1)
    mask = target != void_code
    return (input.argmax(dim=1)[mask]==target[mask]).float().mean()

metrics = drone_accuracy_mask
wd=1e-2    # wd = weight decay

Fastai’s unet_learner

  • このモジュールは、ImageNetで事前トレーニングされたバックボーンから動的U-Netを構築し、中間サイズを自動的に推測します。
  • 以下の図がオリジナルのU-Netです。 ここでの違いは、左側が事前学習済みモデルであることです。
  • このU-Netは、エンコーダー(resnet50などの事前トレーニング済みモデルである場合もあります)の上に配置され、最終的な出力はnum_classesになります。 semantic-segmentation


arch = models.resnet34
learn = unet_learner(data, 
                     arch, 
                     metrics = [metrics], 
                     wd = wd, bottle=True, 
                     model_dir = '/kaggle/working/') # モデルを保存するためのディレクトリを指定

モデルのSumarry

print(learn.summary())

semantic-segmentation

(略)

モデルに適した学習率を探す

learn.lr_find()
learn.recorder.plot()

semantic-segmentation


gc.collect() # to clear the cache

callbacks = SaveModelCallback(learn, monitor = 'drone_accuracy_mask', every = 'improvement', mode='max', name = 'best_model' )
lr = 1e-3           # Learning Rate
learn.fit_one_cycle(20, slice(lr), pct_start = 0.8,callbacks = [callbacks])

semantic-segmentation


learn.lr_find()
learn.recorder.plot()

semantic-segmentation


gc.collect()

callbacks2 = SaveModelCallback(learn, monitor = 'drone_accuracy_mask', every = 'improvement', mode='max', name = 'best_model_ft')
learn.fit_one_cycle(10, slice(1e-6), pct_start = 0.8,callbacks = [callbacks2])

semantic-segmentation


Intial dynamic unet on top of an encoder ( resnet34 pretrained = 'imagenet' ), trained for 30 epochs gave an accuracy of 80.00% .

学習したモデルを確認

learn.show_results(rows=2, figsize=(16,9))

semantic-segmentation


learn.save('stage-1-big')  # saving the model 

モデルをエクスポート

learn.export('/kaggle/working/drone_mask.pkl')

予測

予測する関数

def drone_predict(f):
    img = open_image(f).resize((3,200,300))
    mask = learn.predict(img)[0]
    _,axs = plt.subplots(1,3, figsize=(24,10))
    img.show(ax=axs[0], title='no mask')
    img.show(ax=axs[1], y=mask, title='masked')
    mask.show(ax=axs[2], title='mask only', alpha=1.)


予測実行

n = random.randint(20,200)
drone_predict(fnames[n])

semantic-segmentation


参考文献