オフロード画像のセグメンテーションチャレンジ チュートリアル

create date : Jan. 8, 2021 at 18:34:25
こちらは、コンペティション「オフロード画像のセグメンテーションチャレンジ」のチュートリアルです。

概要

こちらのチュートリアルでは、簡易的なモデリングを用いて予測値を算出し、提出物を作成する方法について説明します。
説明は以下の流れで構成されています。

  1. 準備
    • 分析環境
    • ディレクトリ構成
    • ライブラリのインポート
  2. データの読み込みと確認
    • データの読み込み
    • 画像データの可視化
    • 各評価対象カテゴリが含まれる画像枚数の確認
  3. モデリング
    • 前処理クラスの定義
    • データセットの作成
    • データローダーの作成
    • モデルの初期化
    • 損失関数、最適化手法の定義
    • モデルの学習を実行する関数の定義
    • 学習の実行
    • 推論結果の可視化
  4. 提出物の作成
    • A. 認識精度部門
    • B. 推論速度部門
  5. 評価スコアの改善に向けて

0. 準備

分析環境

当ファイルのソースコードは、以下の環境で正常に動作することを確認しています。

  • OS: Ubuntu 16.04.7
  • 言語: Python==3.7.9
  • ライブラリ:
    • numpy==1.19.1
    • pandas==1.1.1
    • matplotlib==3.3.1
    • torch==1.6.0
    • torchvision==0.7.0

ディレクトリ構成

当ファイルは以下のディレクトリ構成のもと実行されることを想定しています。

  • off_road/
    • train_images_A/
      • train_image_A0000.png
      • train_image_A0001.png
      • ...
      • train_image_A3006.png
    • train_annotations_A/
      • train_annotation_A0000.png
      • train_annotation_A0001.png
      • ...
      • train_annotation_A3006.png
    • precision_test_images/
      • precision_test_image_0000.png
      • precision_test_image_0001.png
      • ...
      • precision_test_image_0639.png
    • 当チュートリアルの実行ファイル

ライブラリのインポート

必要なライブラリのインポートを行います。

In [1]:
import os
from glob import glob
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['font.sans-serif'] = ['Hiragino Maru Gothic Pro', 'Yu Gothic', 'Meirio', 'Takao', 'IPAexGothic', 'IPAPGothic', 'VL PGothic', 'Noto Sans CJK JP']
from PIL import Image
import torch
import torch.utils.data as data
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms

print(torch.__version__)
print(torch.cuda.is_available())

# 評価対象カテゴリ
eval_names = ('road','dirt road', 'other obstacle')
eval_colors = ((128, 64, 128), (255, 128, 128), (0, 0, 70))
1.6.0
True

1. データの読み込み

データの読み込みを行います。
まずは「学習用画像」「学習用アノテーション」「精度評価用画像」のファイルパスをそれぞれ取得しましょう。

In [2]:
train_images_path_list = sorted(glob('off_road/train_images_A/*.png'))
train_annotations_path_list = sorted(glob('off_road/train_annotations_A/*.png'))
precision_test_images_path_list = sorted(glob('off_road/precision_test_images/*.png'))

print('================')
print('学習用画像: ')
print(len(train_images_path_list))
print(train_images_path_list[:5])
print('================')
print('学習用アノテーション: ')
print(len(train_annotations_path_list))
print(train_annotations_path_list[:5])
print('================')
print('精度評価用画像: ')
print(len(precision_test_images_path_list))
print(precision_test_images_path_list[:5])
================
学習用画像: 
3007
['off_road/train_images_A/train_image_A0000.png', 'off_road/train_images_A/train_image_A0001.png', 'off_road/train_images_A/train_image_A0002.png', 'off_road/train_images_A/train_image_A0003.png', 'off_road/train_images_A/train_image_A0004.png']
================
学習用アノテーション: 
3007
['off_road/train_annotations_A/train_annotation_A0000.png', 'off_road/train_annotations_A/train_annotation_A0001.png', 'off_road/train_annotations_A/train_annotation_A0002.png', 'off_road/train_annotations_A/train_annotation_A0003.png', 'off_road/train_annotations_A/train_annotation_A0004.png']
================
精度評価用画像: 
640
['off_road/precision_test_images/precision_test_image_0000.png', 'off_road/precision_test_images/precision_test_image_0001.png', 'off_road/precision_test_images/precision_test_image_0002.png', 'off_road/precision_test_images/precision_test_image_0003.png', 'off_road/precision_test_images/precision_test_image_0004.png']

学習用データは、「画像」・「アノテーション」ともに3007枚、
精度評価用の画像は640枚存在することを確認しました。

画像データの可視化

画像データの可視化を実行します。
上で取得したファイルパスを用いて画像データを読み込み、互いに対応する「画像」と「アノテーション」を横に並べて可視化してみましょう。

In [3]:
# 画像の読み込み
image_0000 = Image.open('off_road/train_images_A/train_image_A0000.png')
annotation_0000 = Image.open('off_road/train_annotations_A/train_annotation_A0000.png')

# 可視化
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 10))

axes[0].imshow(image_0000)
axes[0].set_title('train_image_A0000.png')

axes[1].imshow(annotation_0000)
axes[1].set_title('train_annotation_A0000.png')

plt.show()

各評価対象カテゴリが含まれる画像枚数の確認

各評価対象カテゴリに該当する物体(road, dirt road, other obstacle)は、全ての画像内に登場するとは限りません。
各カテゴリの物体が各画像にどれほどの頻度で登場しているのかについて、学習用アノテーション画像を対象に確認しましょう。

In [4]:
count = {
    'road': 0,
    'dirt road': 0,
    'other obstacle':0
}

for train_annotation_path in train_annotations_path_list:
    image = np.array(Image.open(train_annotation_path))
    
    for eval_name, eval_color in zip(eval_names, eval_colors):
        mask = (image==eval_color).sum(axis=2)==3
        if np.any(mask):
            count[eval_name] += 1
            
plt.bar(count.keys(), count.values())
Out[4]:

「dirt road」「other obstacle」と比べて「road」の出現頻度は低いようです。

2. モデリング

簡易的な手法を用いたモデリングを実施します。

前処理クラスの定義

前処理では、「画像の縮小」「テンソル化」「入力画像の標準化」「アノテーション画像の4カテゴリ表現への変換」などを行います。

In [5]:
class OffRoadTransform():
    def __init__(self, image_size, mean, std):
        self.image_size = image_size
        self.mean = mean
        self.std = std
        
    def __call__(self, image, annotation):
        
        # リサイズ
        image = image.resize((self.image_size[1], self.image_size[0]))        
        annotation = annotation.resize((self.image_size[1], self.image_size[0]))
        
        # テンソル化&標準化
        image = transforms.functional.to_tensor(image)
        image = transforms.functional.normalize(image, self.mean, self.std)
        
        # アノテーション画像の色(RGB)情報を以下のように対応するようマッピングし、2次元の配列に変換する
        """
        road(128, 64, 128) -> 1
        dirt road(255, 128, 128) -> 2
        other obstacle(0, 0, 70) -> 3
        上記以外 -> 0
        """
        annotation = np.array(annotation)
        converted_annotation = np.zeros(annotation.shape[:-1])
        for i, eval_color in enumerate(eval_colors):
            mask = (annotation==eval_color).sum(axis=2)==3
            converted_annotation[mask] = i+1
        annotation = torch.from_numpy(converted_annotation)
        
        return image, annotation

データセットの作成

torch.utils.data.Datasetクラスを継承したクラスを作成します。

In [6]:
class OffRoadDataset(data.Dataset):
    def __init__(self, image_list, annotation_list, transform):
        self.image_list = image_list
        self.annotation_list = annotation_list
        self.transform = transform

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, index):

        image_filepath = self.image_list[index]
        annotation_filepath = self.annotation_list[index]
        
        image = Image.open(image_filepath)
        annotation = Image.open(annotation_filepath)

        image, annotation = self.transform(image, annotation)
        
        return image, annotation
In [7]:
train_dataset = OffRoadDataset(train_images_path_list, train_annotations_path_list,
                              transform=OffRoadTransform(image_size=(270, 480), mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)))

print(train_dataset.__getitem__(0)[0].shape)
print(train_dataset.__getitem__(0)[1].shape)
torch.Size([3, 270, 480])
torch.Size([270, 480])

データローダーの作成

In [8]:
batch_size = 8
train_dataloader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

モデルの初期化

モデルの初期化を行います。
ここでは、torchvisionライブラリに標準で実装されているdeeplabv3_resnet101モデルを使用します。

In [9]:
net = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)

当コンペティションの課題内容に合わせて、モデルの出力層のチャンネル数を変更しましょう。
3つの評価対象カテゴリ(「road」「dirt road」「other obstacle」)及び「その他」の合計4つのカテゴリに分類するということでチャンネル数は4に変更します。

In [10]:
net.classifier[-1] = nn.Conv2d(256, 4, kernel_size=(1, 1), stride=(1, 1))

損失関数、最適化手法の定義

In [11]:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

モデルの学習を実行する関数の定義

モデルの学習を実行する関数を定義しましょう。
全エポックの終了後に学習済みモデルのオブジェクトをpickleファイルとして保存するように設定します。

In [12]:
def train(net, train_dataloader, criterion, optimizer, n_epoch):
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    
    net.to(device)
    net.train()

    for epoch in range(1, n_epoch+1):
        epoch_train_loss = 0.0
        optimizer.zero_grad()
        
        for images, annotations in train_dataloader:
            images = images.to(device)
            annotations = annotations.to(device)
            optimizer.step()
            optimizer.zero_grad()
            
            with torch.set_grad_enabled(True):
                outputs = net(images)['out']
                loss = criterion(outputs, annotations.long())
                loss.backward()
                
        print(f'Epoch {epoch} finished')

    pd.to_pickle(net, "tutorial_model.pkl")

学習の実行

モデルの学習を実行します。

In [13]:
train(net, train_dataloader, criterion, optimizer, n_epoch=5)
cuda:0
Epoch 1 finished
Epoch 2 finished
Epoch 3 finished
Epoch 4 finished
Epoch 5 finished

推論結果の可視化

学習済みのモデルを用いて推論を実行しましょう。
また、元画像と推論の結果生成された画像を並べて可視化することで、モデルの学習が上手くいっているのかを確認してみましょう。

In [14]:
# 学習済みモデルの読み込み
net = pd.read_pickle("tutorial_model.pkl")
In [15]:
# デバイスの設定
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()

# 前処理クラスのインスタンス化
test_transform = OffRoadTransform(image_size=(270, 480), mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))

# 画像データの読み込み
test_image_0000 = Image.open('off_road/precision_test_images/precision_test_image_0000.png')
dummy_annotation = Image.open('off_road/train_annotations_A/train_annotation_A0000.png')

image_transformed, _ = test_transform(test_image_0000, dummy_annotation)
image_transformed = image_transformed.unsqueeze(0)
image_transformed = image_transformed.to(device)

# 推論の実行
prediction = net(image_transformed)['out']
prediction = prediction[0].to('cpu').detach().numpy()
prediction = np.argmax(prediction, axis=0).astype('uint8')
prediction = np.array(Image.fromarray(prediction).resize([1920, 1080]))

# 推論結果をRGB画像に変換
RGB = np.zeros([1080, 1920, 3], dtype='uint8')
for i, color in enumerate(eval_colors):
    mask = prediction==i+1
    RGB[mask] = color
RGB_prediction = np.array(Image.fromarray(RGB).resize([1920, 1080]))

# 可視化
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 10))

axes[0].imshow(test_image_0000)
axes[0].set_title('元画像')

axes[1].imshow(RGB_prediction)
axes[1].set_title('予測結果')
Out[15]:
Text(0.5, 1.0, '予測結果')