セマンティックセグメンテーションを試してみる(Pytorch)

##初めに

セマンティックセグメンテーションは画像認識技術の一種で、画素毎に認識することができます。
<img width="418" alt="seg.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/1976e9e0-635f-1e00-e153-76eaca63e5d5.png">

詳しい理論等は別に譲りますが、Pytorchを用いてセマンティックセグメンテーションを試してみたいと思います。今回はSeg-NetとかU-netとかPSP-netのような層が深くて複雑な構造のネットワークではなく、もっと浅くてシンプルで、ノートPCでも十分学習可能なネットワークを扱います。

環境は
CPU: intel(R) core(TM)i5 7200U
メモリ: 8 GB
OS: Windows10
python ver3.6.9
pytorch ver1.3.1
numpy ver1.17.4


## データセットの作成
今回は自分で合成した画像を使います。上の線の画像が入力データ,下の塗りつぶした画像が教師データです。すなわち自動でペイントソフトみたいに塗りつぶしをするネットワークを作成します。
![input_auto.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/75995c67-654e-0680-1828-337f1fca1b60.png)
![correct_auto.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/e110e2ea-22f9-9b3b-ae24-ab736a28859e.png)

学習に必要なデータを作ります。
imgsは入力データ1000枚
imgs_anoが出力データ(教師データ)1000枚
四角と四角は必ず被らないようになっており、辺の長さや四角の数もランダムで決定するようになっています。

```python
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import TensorDataset, DataLoader

def rectangle(img, img_ano, centers, max_side):
    """
    img     …四角形の線のみの2次元画像
    img_ano …あのーテーション画像
    centers …center座標のlist
    max_side…辺の最大長さの1/2 
    """
    if max_side < 3: #max_sideが小さすぎるとき
        max_side = 4
    #辺の長さの1/2を定義
    side_x = np.random.randint(3, int(max_side))
    side_y = np.random.randint(3, int(max_side))    
    
    #中心の座標,(x, y)を定義
    x = np.random.randint(max_side + 1, img.shape[0] - (max_side + 1))
    y = np.random.randint(max_side + 1, img.shape[1] - (max_side + 1))
    
    #過去の中心位置と近い位置が含まれた場合,inputデータをそのまま返す
    for center in centers:
        if np.abs(center[0] - x) < (2 *max_side + 1):
            if np.abs(center[1] - y) < (2 * max_side + 1):
                return img, img_ano, centers
            
    img[x - side_x : x + side_x, y - side_y] = 1.0      #上辺
    img[x - side_x : x + side_x, y + side_y] = 1.0      #下辺
    img[x - side_x, y - side_y : y + side_y] = 1.0      #左辺
    img[x + side_x, y - side_y : y + side_y + 1] = 1.0  #右辺
    img_ano[x - side_x : x + side_x + 1, y - side_y : y + side_y + 1] = 1.0
    centers.append([x, y])
    return img, img_ano, centers


num_images = 1000                                   #生成する画像数
length = 64                                          #画像のサイズ
imgs = np.zeros([num_images, 1, length, length])     #ゼロ行列を生成,入力画像
imgs_ano = np.zeros([num_images, 1, length, length]) #出力画像

for i in range(num_images):
    centers = []
    img = np.zeros([length, length])
    img_ano = np.zeros([64, 64])
    for j in range(6):                       #四角形を最大6つ生成
        img, img_ano, centers = rectangle(img, img_ano, centers, 12) 
    imgs[i, 0, :, :] = img
    imgs_ano[i, 0, :, :] = img_ano
   
imgs = torch.tensor(imgs, dtype = torch.float32)                 #ndarray - torch.tensor
imgs_ano = torch.tensor(imgs_ano, dtype = torch.float32)           #ndarray - torch.tensor
data_set = TensorDataset(imgs, imgs_ano)
data_loader = DataLoader(data_set, batch_size = 100, shuffle = True)
```

###ネットワーク_1

次にPytorchでネットワークのクラスを作成します。
まずは[前回](https://qiita.com/shushin/items/3f35a9a4200d7be74dc9)のオートエンコーダで定義したネットワークをそのまま使いました。オートエンコーダもセグメンテーションも入力画像と同じサイズの画像を生成するので、(Pytorchの場合)使えました。Tensorflowとかどうなっているのだろう?

```python
class ConvAutoencoder(nn.Module):
    def __init__(self):
        super(ConvAutoencoder, self).__init__()
        #Encoder Layers
        self.conv1 = nn.Conv2d(in_channels = 1,
                               out_channels = 16,
                               kernel_size = 3,
                               padding = 1)
        self.conv2 = nn.Conv2d(in_channels = 16,
                               out_channels = 4,
                               kernel_size = 3,
                               padding = 1)
        #Decoder Layers
        self.t_conv1 = nn.ConvTranspose2d(in_channels = 4, out_channels = 16,
                                          kernel_size = 2, stride = 2)
        self.t_conv2 = nn.ConvTranspose2d(in_channels = 16, out_channels = 1,
                                          kernel_size = 2, stride = 2)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2, 2)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        #encode#                           
        x = self.relu(self.conv1(x))        
        x = self.pool(x)                  
        x = self.relu(self.conv2(x))      
        x = self.pool(x)                  
        #decode#
        x = self.relu(self.t_conv1(x))    
        x = self.sigmoid(self.t_conv2(x))
        return x
```

このネットワークで学習させます。

```python
#******ネットワークを選択******
net = ConvAutoencoder()                               
loss_fn = nn.MSELoss()                                #損失関数の定義
optimizer = optim.Adam(net.parameters(), lr = 0.01)

losses = []                                     #epoch毎のlossを記録
epoch_time = 30
for epoch in range(epoch_time):
    running_loss = 0.0                          #epoch毎のlossの計算
    net.train()
    for i, (XX, yy) in enumerate(data_loader):
        optimizer.zero_grad()       
        y_pred = net(XX)
        loss = loss_fn(y_pred, yy)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print("epoch:",epoch, " loss:", running_loss/(i + 1))
    losses.append(running_loss/(i + 1))

#lossの可視化
plt.plot(losses)
plt.ylabel("loss")
plt.xlabel("epoch time")
plt.savefig("loss_auto")
plt.show()
```

epoch毎の損失(loss)を可視化したものです。
epochが30回でそれなりに収束した状態でしょうか?
![loss_auto.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/0403c4f2-9f8a-34b4-874d-15d6f22e8648.png)

学習に使っていない画像を用いて、試してみます。
大まかな位置は判別できていますが、境界の付近はうまく取れていない印象があります。
<img width="280" alt="output_auto.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/ae016653-5d3d-1ad9-3dc9-906b08c522b4.png">

```python
net.eval()            #評価モード
#今まで学習していない画像を1つ生成
num_images = 1
img_test = np.zeros([num_images, 1, length, length])
imgs_test_ano = np.zeros([num_images, 1, length, length])
for i in range(num_images):
    centers = []
    img = np.zeros([length, length])
    img_ano = np.zeros([length, length])
    for j in range(6):
        img, img_ano, centers = rectangle(img, img_ano, centers, 7)
    img_test[i, 0, :, :] = img

img_test = img_test.reshape([1, 1, 64, 64])
img_test = torch.tensor(img_test, dtype = torch.float32)
img_test = net(img_test)             #生成した画像を学習済のネットワークへ
img_test = img_test.detach().numpy() #torch.tensor - ndarray
img_test = img_test[0, 0, :, :]

plt.imshow(img, cmap = "gray")       #inputデータの可視化
plt.savefig("input_auto")
plt.show()
plt.imshow(img_test, cmap = "gray")  #outputデータの可視化
plt.savefig("output_auto")
plt.show()
plt.imshow(img_ano, cmap = "gray")   #正解データ
plt.savefig("correct_auto")
plt.plot()
```



###ネットワークを深くしてみる。

先ほどのモデルでは十分な性能を得ることができませんでしたので、層を深くしてみたいと思います。
ここでは単に深くするだけではなく、過学習を防ぐバッチノーマライゼーションと、デコーダでアップサンプリングを入れてみます。アップサンプリングの詳しい解説は[こちらの記事](https://qiita.com/cyberailab/items/d11862852eccc17585e8)が分かりやすかったです。

```python
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        #encoder
        self.encoder_conv_1 = nn.Sequential(*[
                                            nn.Conv2d(in_channels = 1, 
                                                      out_channels = 6,
                                                      kernel_size = 3,
                                                      padding = 1),
                                            nn.BatchNorm2d(6)
                                            ])
        
        self.encoder_conv_2 = nn.Sequential(*[
                                            nn.Conv2d(in_channels = 6,
                                                      out_channels = 16,
                                                      kernel_size = 3,
                                                      padding = 1),
                                            nn.BatchNorm2d(16)
                                            ])
        self.encoder_conv_3 = nn.Sequential(*[
                                            nn.Conv2d(in_channels = 16,
                                                      out_channels = 32,
                                                      kernel_size = 3,
                                                      padding = 1),
                                            nn.BatchNorm2d(32)
                                            ])
        
        #decoder
        self.decoder_convt_3 = nn.Sequential(*[
                                            nn.ConvTranspose2d(in_channels = 32,
                                                               out_channels = 16,
                                                               kernel_size = 3,
                                                               padding = 1),
                                            nn.BatchNorm2d(16)
                                            ])
        
        self.decoder_convt_2 = nn.Sequential(*[
                                            nn.ConvTranspose2d(in_channels = 16,
                                                               out_channels = 6,
                                                               kernel_size = 3,
                                                               padding = 1),
                                            nn.BatchNorm2d(6)
                                            ])
        
        self.decoder_convt_1 = nn.Sequential(*[
                                            nn.ConvTranspose2d(in_channels = 6,
                                                               out_channels = 1,
                                                               kernel_size = 3,
                                                               padding = 1)
                                            ])
    
    def forward(self, x):
        #encoder
        dim_0 = x.size()                    
        x = F.relu(self.encoder_conv_1(x))                            
        x, indices_1 = F.max_pool2d(x, kernel_size = 2,
                                    stride = 2, 
                                    return_indices = True)  #indiceでmaxpoolの位置を記録          
        dim_1 = x.size()
        x = F.relu(self.encoder_conv_2(x))                            
        x, indices_2 = F.max_pool2d(x, kernel_size = 2,
                                    stride = 2, 
                                    return_indices = True)            
        
        dim_2 = x.size()
        x = F.relu(self.encoder_conv_3(x))
        x, indices_3 = F.max_pool2d(x, kernel_size = 2,
                                    stride = 2, 
                                    return_indices = True)
        
        #decoder
        x = F.max_unpool2d(x, indices_3, kernel_size = 2,
                           stride = 2, output_size = dim_2)
        x = F.relu(self.decoder_convt_3(x))
        
        x = F.max_unpool2d(x, indices_2, kernel_size = 2,
                           stride = 2, output_size = dim_1)           
        x = F.relu(self.decoder_convt_2(x))                           
        
        x = F.max_unpool2d(x, indices_1, kernel_size = 2,
                           stride = 2, output_size = dim_0)           
        x = F.relu(self.decoder_convt_1(x))                           
        x = torch.sigmoid(x)                                       
        
        return x
```

このネットワークに切り替えるのは簡単で

```python
#******ネットワークを選択******
net = ConvAutoencoder()
```
となっている場所を、新たに作成したクラスに変更するだけです。

```python
#******ネットワークを選択******
net = Net()
```

lossの変遷をグラフ化します。

![loss_auto.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/3e0ea0cb-dd4c-6f10-2a50-1567accf180d.png)

学習に使用していないデータを入力して、正解画像と比較してみます。
<img width="425" alt="output.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/141993/a0c0a665-cb5a-be82-dba1-56845116803f.png">

セグメンテーションできていることが分かります。


## 終わりに
簡単なセグメンテーションを今回試しました。実用には程遠いシンプルなモデルでしたが、雰囲気は大体つかめたような気がします。