SegNetにスキップ構造を組み込んで性能比較(CaDIS: a Cataract Dataset)

# 1. はじめに
　本記事では[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)行ったCaDIS: a Cataract Dataset[^1]を対象として、SegNet[^2]を用いてセマンティックセグメンテーションした記事の続編です。今回は[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)のネットワークにスキップ構造を組み込んだU-Net[^3]性能比較を行います。

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

# 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. データセット&データ分割
　[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)と同じなのでこちらを参照してください。

# 4. モデル構築&学習
　今回は[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)のSegNet[^2]にスキップ構造を組み込んだU-Net[^3]を実装し、性能比較を行います。U-Netは下図のようなネットワークでエンコーダーMax Pooling前の出力をデコーダーのUp Sampling後の出力を連結させてConvolution層に入力します。これにより、Max Poolingで画像を圧縮することにより消えてしまった情報を回復することができるようになっています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/466418e9-a896-9775-0c11-d63fbf992891.png)

　今回は[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)のネットワーク構造にスキップ構造だけを導入します。まず、ネットワークの分岐は以下のように行います。

```py
x = Conv2D(32, (3, 3), padding='same', kernel_initializer='he_normal')(inputs)
x = BatchNormalization()(x)
x1 = Activation('relu')(x)
x = MaxPool2D(pool_size=(2, 2))(x1)
```

`x1`を更新せずに残しておくことで出力の分岐を行います。次に分岐した出力の連結は以下のように行います。

```py
x = UpSampling2D(size=(2, 2))(x)
x = concatenate([x, x1], axis=-1)
x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
```

`concatenate()`を使用すること入力テンソルの連結が行えます。今回は`x1=(height, width, ch1)`と`x=(height, width, ch2)`を結合して`(height, width, ch1+ch2)`の出力テンソルを得ます。これを各Max Pooling、Up Samplingで適用したネットワークを以下に示します。

<details><summary>U-Netのコード</summary><div>

```py
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, MaxPool2D, UpSampling2D, concatenate
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation

# U-Net(エンコーダー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)
    x1 = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x1)

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

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

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

    ## 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)
    x5 = Activation('relu')(x)
    x = MaxPool2D(pool_size=(2, 2))(x5)

    ## 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 = concatenate([x, x5], axis=-1)
    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 = concatenate([x, x4], axis=-1)
    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 = concatenate([x, x3], axis=-1)
    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 = concatenate([x, x2], axis=-1)
    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 = concatenate([x, x1], axis=-1)
    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)
```
</div></details>

さらに、今後深層化などによりネットワーク部分の記述が冗長になってくると思うので、エンコーダーの`Conv+BN+Relu+MaxPool`とデコーダーの`Conv+BN+Relu+Upsampl`を関数化します。関数化してちょっとだけスッキリしたものを以下に示します。

```py
import dataclasses
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, MaxPool2D, UpSampling2D, concatenate
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation

# U-Net(エンコーダー8層、デコーダー8層)の構築
@dataclasses.dataclass
class CNN:
    input_shape: tuple # 入力画像サイズ
    classes: int # 分類クラス数

    def __post_init__(self):
        # 入力画像サイズは32の倍数でなければならない
        assert self.input_shape[0]%32 == 0, 'Input size must be a multiple of 32.'
        assert self.input_shape[1]%32 == 0, 'Input size must be a multiple of 32.'


    # エンコーダーブロック
    @staticmethod
    def encoder(x, blocks, filters, pooling):
        for i in range(blocks):
            x = Conv2D(filters, (3, 3), padding='same', kernel_initializer='he_normal')(x)
            x = BatchNormalization()(x)
            x = Activation('relu')(x)

        if pooling:
            return MaxPool2D(pool_size=(2, 2))(x), x
        else:
            return x


    # デコーダーブロック
    @staticmethod
    def decoder(x1, x2, blocks, filters):
        x = UpSampling2D(size=(2, 2))(x1)
        x = concatenate([x, x2], axis=-1)

        for i in range(blocks):
            x = Conv2D(filters, (3, 3), padding='same', kernel_initializer='he_normal')(x)
            x = BatchNormalization()(x)
            x = Activation('relu')(x)

        return x


    def create(self):
        # エンコーダー
        inputs = Input(shape=(self.input_shape[0], self.input_shape[1], 3)) # 入力層
        x, x1 = self.encoder(inputs, blocks=1, filters=32, pooling=True) # 1層目
        x, x2 = self.encoder(x, blocks=1, filters=64, pooling=True) # 2層目
        x, x3 = self.encoder(x, blocks=1, filters=128, pooling=True) # 3層目
        x, x4 = self.encoder(x, blocks=1, filters=256, pooling=True) # 4層目
        x, x5 = self.encoder(x, blocks=2, filters=512, pooling=True) # 5、6層目
        x = self.encoder(x, blocks=2, filters=1024, pooling=False) # 7、8層目

        # デコーダー
        x = self.encoder(x, blocks=1, filters=1024, pooling=False) # 1層目
        x = self.decoder(x, x5, blocks=2, filters=512) # 2、3層目
        x = self.decoder(x, x4, blocks=1, filters=256) # 4層目
        x = self.decoder(x, x3, blocks=1, filters=128) # 5層目
        x = self.decoder(x, x2, blocks=1, filters=64) # 6層目
        ## 7、8層目
        x = UpSampling2D(size=(2, 2))(x)
        x = concatenate([x, x1], axis=-1)
        x = Conv2D(64, (3, 3), strides=(1, 1), padding='same', kernel_initializer='he_normal')(x)
        x = Conv2D(self.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(input_shape=image_size, classes=classes).create()
```

学習方法やパラメータなどは[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)と同じです。学習結果は以下のようになりました。
![history.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405329/21074564-d6ad-c417-263a-8d1e0cb50cd9.jpeg)


# 5. 評価
　評価は[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)と同じようにクラスごとのaverage IoUとそれの平均を取ったmean IoUで行います。

以下、評価結果です。表から、[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)は0%だった多くのクラスが若干ですが推論できるようになっていることがわかります。また、mean IoUは18.5%となりました。[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)のSegNet[^2]は15.0%だったので3.5%精度向上させることができました。

|Index|Class|average IoU(SegNet)[%]|average IoU(U-Net)[%]|
|:--:|:--:|:--:|:--:|
|0|Pupil|85.3|86.5|
|1|Surgical Tape|53.3|57.1|
|2|Hand|6.57|6.96|
|3|Eye Retractors|21.9|53.6|
|4|Iris|74.4|76.0|
|5|Eyelid|0|0|
|6|Skin|49.7|48.4|
|7|Cornea|88.0|88.5|
|8|Hydro. Cannula|0|31.8|
|9|Visco. Cannula|0|4.36|
|10|Cap. Cystotome|0|3.71|
|11|Rycroft Cannula|0|4.37|
|12|Bonn Forceps|3.58|7.94|
|13|Primary Knife|5.35|10.3|
|14|Phaco. Handpiece|0.0781|12.3|
|15|Lens Injector|16.4|15.8|
|16|A/I Handpiece|16.4|20.5|
|17|Secondary Knife|6.08|11.8|
|18|Micromanipulator|0|8.99|
|19|A/I Handpiece Handle|6.49|8.16|
|20|Cap. Forceps|0|0.337|
|21|Rycroft Cannula Handle|0|0.00863|
|22|Phaco. Handpiece Handle|0|4.26|
|23|Cap. Cystotome Handle|0|0.407|
|24|Secondary Knife Handle|2.49|3.82|
|25|Lens Injector Handle|0|0|
|26|Water Sprayer|─|─|
|27|Suture Needle|0|0|
|28|Needle Holder|─|─|
|29|Charleux Cannula|0|0|
|30|Vannas Scissors|─|─|
|31|Primary Knife Handle|0|0|
|32|Viter. Handpiece|0|0|
|33|Mendez Ring|─|─|
|34|Biomarker|─|─|
|35|Marker|─|─|

# 7. まとめ
　本記事では、[前回](https://qiita.com/burokoron/items/c730e607607c925c6fd1)のSegNet[^2]にスキップ構造を組み込んだU-Net[^3]の実装を行いました。CaDIS: a Cataract Dataset[^1]を対象として、mean IoUで性能比較を行い、3.5%の精度向上を確認しました。論文[^1]の最高精度はPSPNetの52.66%まであと34.16%となりました。今後もこの結果をベースとしてネットワーク構造やデータ拡張方法などの最新の手法を取り入れながら、同等以上の性能を目指そうと思います。

[^1]: [CaDIS: Cataract Dataset for Image Segmentation] (https://arxiv.org/abs/1906.11586)
[^2]: [SegNet] (http://mi.eng.cam.ac.uk/projects/segnet/)
[^3]: [U-Net: Convolutional Networks for Biomedical Image Segmentation] (https://arxiv.org/abs/1505.04597)