CaDIS: a Cataract Datasetで画像セグメンテーション

# 1. はじめに
　本記事では入門向けとして、TensorFlow2.0を使用して、とりあえずDeep Learningでセマンティック・セグメンテーションを行うことを目標とします。画像データセットはDigital Surgery Ltd.により公開されている白内障手術セグメンテーションデータセット[^1]を使用します。また、ネットワークは[前回] (https://qiita.com/burokoron/items/8c011c219b7545c50355)使用した10層CNNをエンコーダとしたSegNet[^2]とします。

[すべてのコード](https://github.com/burokoron/StaDeep/tree/master/SegNet)

# 2. 環境
- PCスペック
  - CPU: Intel Core i9-9900K
  - RAM: 16GB
  - GPU: NVIDIA GeForce GTX 1080 Ti
- ライブラリ
  - Python: 3.7.4
  - numpy: 1.16.5
  - matplotlib: 3.1.1
  - opencv: 3.4.1
  - pandas: 0.25.1
  - tqdm: 4.31.1
  - scikit-learn: 0.21.3
  - tensorflow-gpu: 2.0.0

# 3. CaDIS: Cataract Dataset for Image Segmentation
　Digital Surgery Ltd.により公開された4738枚(25動画)の白内障手術セグメンテーションデータセット。下記リンクより手術画像とセグメンテーション画像をダウンロードできます。
CaDIS Dataset
https://cataracts.grand-challenge.org/CaDIS/

セグメンテーションラベルは以下の通りです。同時に各クラスが占めるピクセル単位での比率も示します。表より、それぞれのグループで存在しないクラスがあることがわかります。~~めんどくさいな！~~学習、検証、テストの画像枚数はそれぞれ3584枚(19動画)、540枚(3動画)、614枚(3動画)です。

|Index|Class|ピクセル比率(学習)[%]|ピクセル比率(検証)[%]|ピクセル比率(テスト)[%]|
|:--:|:--:|:--:|:--:|:--:|
|0|Pupil|17.1|15.7|16.2|
|1|Surgical Tape|6.51|6.77|4.81|
|2|Hand|0.813|0.725|0.414|
|3|Eye Retractors|0.564|0.818|0.388|
|4|Iris|11.0|11.0|12.8|
|5|Eyelid|0|0|1.86|
|6|Skin|12.0|20.4|10.7|
|7|Cornea|49.6|42.2|50.6|
|8|Hydro. Cannula|0.138|0.0984|0.0852|
|9|Visco. Cannula|0.0942|0.0720|0.0917|
|10|Cap. Cystotome|0.0937|0.0821|0.0771|
|11|Rycroft Cannula|0.0618|0.0788|0.0585|
|12|Bonn Forceps|0.241|0.161|0.276|
|13|Primary Knife|0.123|0.258|0.249|
|14|Phaco. Handpiece|0.173|0.240|0.184|
|15|Lens Injector|0.343|0.546|0.280|
|16|A/I Handpiece|0.327|0.380|0.305|
|17|Secondary Knife|0.102|0.0933|0.148|
|18|Micromanipulator|0.188|0.229|0.215|
|19|A/I Handpiece Handle|0.0589|0.0271|0.0358|
|20|Cap. Forceps|0.0729|0.0144|0.0384|
|21|Rycroft Cannula Handle|0.0406|0.0361|0.0101|
|22|Phaco. Handpiece Handle|0.0566|0.00960|0.0202|
|23|Cap. Cystotome Handle|0.0170|0.0124|0.0287|
|24|Secondary Knife Handle|0.0609|0.0534|0.0124|
|25|Lens Injector Handle|0.0225|0.0599|0.0382|
|26|Water Sprayer|0.000448|0|0.00361|
|27|Suture Needle|0.000764|0|0|
|28|Needle Holder|0.0201|0|0|
|29|Charleux Cannula|0.00253|0|0.0164|
|30|Vannas Scissors|0.00107|0|0|
|31|Primary Knife Handle|0.000321|0|0.000385|
|32|Viter. Handpiece|0|0|0.0782|
|33|Mendez Ring|0.0960|0|0|
|34|Biomarker|0.00619|0|0|
|35|Marker|0.0661|0|0|

また、以下に画像サンプルを示します。生のセグメンテーション画像は上表のIndexをそのまま画素値としたグレースケール画像となっています。

<details><summary>グロテスクな画像が含まれます</summary><div>
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/1b5f0367-79d1-2e9c-7287-24e8f3a6fca1.png" width=40%> <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/b1aca52b-b172-56ac-1682-3e5c8516f8ab.png" width=40%>
手術画像とセグメンテーション画像[^1]
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/c4a2a9f3-8d64-6fcc-4cf4-a57561e12cd4.png" width=40%>
生のセグメンテーション画像
</div></details>

# 4. データ分割
 このデータセットでは学習、検証、テストに使用すべき画像(動画)が決められています。データセット内のsplits.txtというファイルに詳細が記述されています。したがって、分割グループはsplits.txtの内容を採用し、各グループの手術画像とセグメンテーション画像のファイルパスおよびその対応を以下のコードでcsvファイルに記述します。

<details><summary>画像ファイルパスをcsvファイルに記述するコード</summary><div>

```py
import os
from collections import defaultdict
import pandas as pd


# 画像とラベルの対応を記述したcsvファイルを作成する
def make_csv(fpath, dirlist):
    # 学習画像のファイルパスを調べる
    dataset = defaultdict(list)
    for dir in dirlist:
        filelist = sorted(os.listdir(f'CaDIS/{dir}/Images'))
        dataset['filename'] += list(map(lambda x: f'{dir}/Images/{x}', filelist))
        filelist = sorted(os.listdir(f'CaDIS/{dir}/Labels'))
        dataset['label'] += list(map(lambda x: f'{dir}/Labels/{x}', filelist))

    # csvファイルで保存
    dataset = pd.DataFrame(dataset)
    dataset.to_csv(fpath, index=False)



# 学習データのビデオフォルダ
train_dir = ['Video01', 'Video03', 'Video04', 'Video06', 'Video08', 'Video09',
             'Video10', 'Video11', 'Video13', 'Video14', 'Video15', 'Video17',
             'Video18', 'Video20', 'Video21', 'Video22', 'Video23', 'Video24',
             'Video25']

# 検証データのビデオフォルダ
val_dir = ['Video05', 'Video07', 'Video16']

# テストデータのビデオフォルダ
test_dir = ['Video02', 'Video12', 'Video19']


# 学習データの画像とラベルの対応を記述したcsvファイルを作成する
make_csv('train.csv', train_dir)

# 検証データの画像とラベルの対応を記述したcsvファイルを作成する
make_csv('val.csv', val_dir)

# 学習データの画像とラベルの対応を記述したcsvファイルを作成する
make_csv('test.csv', test_dir)
```
</div></details>

学習、検証、テストデータのファイルパスが書かれたcsvファイルはこのような形式にしました。

|filename|label|
|:--:|:--:|
|Video01/Images/Video1_frame000090.png|Video01/Labels/Video1_frame000090.png|
|Video01/Images/Video1_frame000100.png|Video01/Labels/Video1_frame000100.png|
|Video01/Images/Video1_frame000110.png|Video01/Labels/Video1_frame000110.png|

# 5. モデル構築&学習
 まずは使用するライブラリをインポートします。

```py
import dataclasses
import math
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, MaxPool2D, UpSampling2D
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import categorical_crossentropy
from tensorflow.keras.utils import Sequence
import cv2
```

次にパラメータ等を記述します。

```py
directory = 'CaDIS' # 画像が保存されているフォルダ
df_train = pd.read_csv('train.csv') # 学習データの情報がかかれたDataFrame
df_validation = pd.read_csv('val.csv') # 検証データの情報がかかれたDataFrame
image_size = (224, 224) # 入力画像サイズ
classes = 36 # 分類クラス数
batch_size = 32 # バッチサイズ
epochs = 300 # エポック数
loss = cce_dice_loss # 損失関数
optimizer = Adam(lr=0.001, amsgrad=True) # 最適化関数
metrics = dice_coeff # 評価方法
# ImageDataGenerator画像増幅のパラメータ
aug_params = {'rotation_range': 5,
              'width_shift_range': 0.05,
              'height_shift_range': 0.05,
              'shear_range': 0.1,
              'zoom_range': 0.05,
              'horizontal_flip': True,
              'vertical_flip': True}
```

学習時のコールバック処理として以下を適用します。

```py
# val_lossが最小になったときのみmodelを保存
mc_cb = ModelCheckpoint('model_weights.h5',
                        monitor='val_loss', verbose=1,
                        save_best_only=True, mode='min')
# 学習が停滞したとき、学習率を0.2倍に
rl_cb = ReduceLROnPlateau(monitor='loss', factor=0.2, patience=3,
                          verbose=1, mode='auto',
                          min_delta=0.0001, cooldown=0, min_lr=0)
# 学習が進まなくなったら、強制的に学習終了
es_cb = EarlyStopping(monitor='loss', min_delta=0,
                      patience=5, verbose=1, mode='auto')
```

学習データと検証データのジェネレータを生成します。データ拡張には`ImageDataGenerator`を使用します。また、今回は`Sequence`を利用してミニバッチデータを作成します。

`__getitem__`関数が具体的にミニバッチを作成する部分です。入力画像の処理は以下の手順で行っています。

1. 画像を読み込む
2. 指定した入力画像サイズにリサイズする
3. float型に変換する
4. データ拡張処理をする
5. 値を255で割り、0～1に正規化する

セグメンテーション画像の処理は以下の手順で行っています。

1. 画像を読み込む
2. 指定した入力画像サイズにリサイズする
3. float型に変換する
4. データ拡張処理をする
5. クラス0の画素は1、そうでなければ0である画像、クラス1の画素は1、そうでなければ0となる画像……(クラス数分)を作成し、チャンネル方向に連結させることで、(縦、横、クラス数)のサイズの配列を作成する

```py
# データのジェネレータ
@dataclasses.dataclass
class TrainSequence(Sequence):
    directory: str # 画像が保存されているフォルダ
    df: pd.DataFrame # データの情報がかかれたDataFrame
    image_size: tuple # 入力画像サイズ
    classes: int # 分類クラス数
    batch_size: int # バッチサイズ
    aug_params: dict # ImageDataGenerator画像増幅のパラメータ

    def __post_init__(self):
        self.df_index = list(self.df.index)
        self.train_datagen = ImageDataGenerator(**self.aug_params)

    def __len__(self):
        return math.ceil(len(self.df_index) / self.batch_size)

    def __getitem__(self, idx):
        batch_x = self.df_index[idx * self.batch_size:(idx+1) * self.batch_size]

        x = []
        y = []
        for i in batch_x:
            rand = np.random.randint(0, int(1e9))
            # 入力画像
            img = cv2.imread(f'{self.directory}/{self.df.at[i, "filename"]}')
            img = cv2.resize(img, self.image_size, interpolation=cv2.INTER_LANCZOS4)
            img = np.array(img, dtype=np.float32)
            img = self.train_datagen.random_transform(img, seed=rand)
            img *= 1./255
            x.append(img)

            # セグメンテーション画像
            img = cv2.imread(f'{self.directory}/{self.df.at[i, "label"]}', cv2.IMREAD_GRAYSCALE)
            img = cv2.resize(img, self.image_size, interpolation=cv2.INTER_LANCZOS4)
            img = np.array(img, dtype=np.float32)
            img = np.reshape(img, (self.image_size[0], self.image_size[1], 1))
            img = self.train_datagen.random_transform(img, seed=rand)
            img = np.reshape(img, (self.image_size[0], self.image_size[1]))
            seg = []
            for label in range(self.classes):
                seg.append(img == label)
            seg = np.array(seg, np.float32)
            seg = seg.transpose(1, 2, 0)
            y.append(seg)

        x = np.array(x)
        y = np.array(y)


        return x, y

# ジェネレータの生成
## 学習データのジェネレータ
train_generator = TrainSequence(directory=directory, df=df_train,
                                image_size=image_size, classes=classes,
                                batch_size=batch_size, aug_params=aug_params)
step_size_train = len(train_generator)
## 検証データのジェネレータ
validation_generator = TrainSequence(directory=directory, df=df_validation,
                                     image_size=image_size, classes=classes,
                                     batch_size=batch_size, aug_params={})
step_size_validation = len(validation_generator)
```

[前回](https://qiita.com/burokoron/items/8c011c219b7545c50355)作成した10層の単純なCNNから全結合を除いた構造をエンコーダー、エンコーダーを逆順にしたような構造をデコーダーとしてSegNetを構築します。SegNetの解説は[こちら](https://qiita.com/cyberailab/items/d11862852eccc17585e8)を参照してください。

```py
# SegNet(エンコーダー8層、デコーダー8層)の構築
def cnn(input_shape, classes):
    # 入力画像サイズは32の倍数でなければならない
    assert input_shape[0]%32 == 0, 'Input size must be a multiple of 32.'
    assert input_shape[1]%32 == 0, 'Input size must be a multiple of 32.'

    # エンコーダー
    ## 入力層
    inputs = Input(shape=(input_shape[0], input_shape[1], 3))

    ## 1層目
    x = Conv2D(32, (3, 3), padding='same', kernel_initializer='he_normal')(inputs)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x)

    ## 2層目
    x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x)

    ## 3層目
    x = Conv2D(128, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x)

    ## 4層目
    x = Conv2D(256, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x)

    ## 5、6層目
    x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x)

    ## 7、8層目
    x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    # デコーダー
    ## 1層目
    x = Conv2D(1024, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    ## 2、3層目
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Conv2D(512, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    ## 4層目
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(256, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    ## 5層目
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(128, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    ## 6層目
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    ## 7、8層目
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = Conv2D(classes, (1, 1), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
    outputs = Activation('softmax')(x)


    return Model(inputs=inputs, outputs=outputs)

# ネットワーク構築
model = cnn(image_size, classes)
model.summary()
model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])
```

あとは[前回](https://qiita.com/burokoron/items/8c011c219b7545c50355)と同じです。学習して、学習曲線を保存します。

```py
# 学習
history = model.fit_generator(
    train_generator, steps_per_epoch=step_size_train,
    epochs=epochs, verbose=1, callbacks=[mc_cb, rl_cb, es_cb],
    validation_data=validation_generator,
    validation_steps=step_size_validation,
    workers=3)

# 学習曲線のグラフを描き保存する
def plot_history(history):
    fig, (axL, axR) = plt.subplots(ncols=2, figsize=(10, 4))

    # [左側] metricsについてのグラフ
    L_title = 'Dice_coeff_vs_Epoch'
    axL.plot(history.history['dice_coeff'])
    axL.plot(history.history['val_dice_coeff'])
    axL.grid(True)
    axL.set_title(L_title)
    axL.set_ylabel('dice_coeff')
    axL.set_xlabel('epoch')
    axL.legend(['train', 'test'], loc='upper left')

    # [右側] lossについてのグラフ
    R_title = "Loss_vs_Epoch"
    axR.plot(history.history['loss'])
    axR.plot(history.history['val_loss'])
    axR.grid(True)
    axR.set_title(R_title)
    axR.set_ylabel('loss')
    axR.set_xlabel('epoch')
    axR.legend(['train', 'test'], loc='upper left')

    # グラフを画像として保存
    fig.savefig('history.jpg')
    plt.close()

# 学習曲線の保存
plot_history(history)
```

学習結果は以下のようになりました。

![history.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/849b2249-d268-e56f-79ae-1b77c48e3167.jpeg)

# 6. 評価
　評価はクラスごとのaverage IoUとそれの平均を取ったmean IoUで行います。計算は以下のコードで行いました。

追加のインポート。

```py
from collections import defaultdict
```

推論と評価は以下の手順で行っています。

1. 画像を読み込む
2. 指定した入力画像サイズにリサイズ
3. float型に変換し、値を0～1に正規化する
4. バッチサイズ1の配列にする
5. 推論し、セグメンテーション画像を取得する
6. セグメンテーション画像サイズを元のサイズにする
7. 各画像、各クラスのIoUを計算する
8. 各クラスのaverage IoUを計算する

```py

    directory = 'CaDIS' # 画像が保存されているフォルダ
    df_test = pd.read_csv('test.csv') # テストデータの情報がかかれたDataFrame
    image_size = (224, 224) # 入力画像サイズ
    classes = 36 # 分類クラス数


    # ネットワーク構築
    model = cnn(image_size, classes)
    model.summary()
    model.load_weights('model_weights.h5')


    # 推論
    dict_iou = defaultdict(list)
    for i in tqdm(range(len(df_test)), desc='predict'):
        img = cv2.imread(f'{directory}/{df_test.at[i, "filename"]}')
        height, width = img.shape[:2]
        img = cv2.resize(img, image_size, interpolation=cv2.INTER_LANCZOS4)
        img = np.array(img, dtype=np.float32)
        img *= 1./255
        img = np.expand_dims(img, axis=0)
        label = cv2.imread(f'{directory}/{df_test.at[i, "label"]}', cv2.IMREAD_GRAYSCALE)

        pred = model.predict(img)[0]
        pred = cv2.resize(pred, (width, height), interpolation=cv2.INTER_LANCZOS4)

        ## IoUの計算
        pred = np.argmax(pred, axis=2)
        for j in range(classes):
            y_pred = np.array(pred == j, dtype=np.int)
            y_true = np.array(label == j, dtype=np.int)
            tp = sum(sum(np.logical_and(y_pred, y_true)))
            other = sum(sum(np.logical_or(y_pred, y_true)))
            if other != 0:
                dict_iou[j].append(tp/other)

    # average IoU
    for i in range(classes):
        if i in dict_iou:
            dict_iou[i] = sum(dict_iou[i]) / len(dict_iou[i])
        else:
            dict_iou[i] = -1
    print('average IoU', dict_iou)
```

以下、評価結果です。また、mean IoUは15.0%となりました。論文[^1]によるとVGGで20.61%なので、こんなもんだと思います。

|Index|Class|average IoU[%]|
|:--:|:--:|:--:|
|0|Pupil|85.3|
|1|Surgical Tape|53.3|
|2|Hand|6.57|
|3|Eye Retractors|21.9|
|4|Iris|74.4|
|5|Eyelid|0.0|
|6|Skin|49.7|
|7|Cornea|88.0|
|8|Hydro. Cannula|0|
|9|Visco. Cannula|0|
|10|Cap. Cystotome|0|
|11|Rycroft Cannula|0|
|12|Bonn Forceps|3.58|
|13|Primary Knife|5.35|
|14|Phaco. Handpiece|0.0781|
|15|Lens Injector|16.4|
|16|A/I Handpiece|16.4|
|17|Secondary Knife|6.08|
|18|Micromanipulator|0|
|19|A/I Handpiece Handle|6.49|
|20|Cap. Forceps|0|
|21|Rycroft Cannula Handle|0|
|22|Phaco. Handpiece Handle|0|
|23|Cap. Cystotome Handle|0|
|24|Secondary Knife Handle|2.49|
|25|Lens Injector Handle|0|
|26|Water Sprayer|─|
|27|Suture Needle|0|
|28|Needle Holder|─|
|29|Charleux Cannula|0|
|30|Vannas Scissors|─|
|31|Primary Knife Handle|0|
|32|Viter. Handpiece|0|
|33|Mendez Ring|─|
|34|Biomarker|─|
|35|Marker|─|

# 7. まとめ
　本記事では、エンコーダー、デコーダーそれぞれ8層のSegNetを用いてDigital Surgery Ltd.により公開された白内障手術セグメンテーションデータセット[^1]のセマンティック・セグメンテーションを行いました。論文[^1]によるとPSPNetだと52.66%出るみたいなので、今後はこの結果をベースとしてネットワーク構造やデータ拡張方法などの最新の手法を取り入れながら、同等以上の性能を目指そうと思います。

[^1]: [CaDIS: Cataract Dataset for Image Segmentation] (https://arxiv.org/abs/1906.11586)
[^2]: [SegNet] (http://mi.eng.cam.ac.uk/projects/segnet/)