U2Netを独自データでトレーニングする

# 好きなオブジェクトを切り抜きたい
既存のセグメンテーションモデルでは、切り抜きたいオブジェクトが対象ではないことがあります。

# 独自のデータでセグメンテーションモデルをトレーニングすればいい

U2netは顕著性セグメンテーションのモデルですが、独自のデータセットでトレーニングして、好きなオブジェクトを切り抜くモデルを作ることができます。
たとえば、（感染症予防の）マスクは、一般的なセグメンテーションデータセットに含まれていませんが、自分でマスクのデータセットを用意してマスクを切り抜くモデルを作成できます。

# 手順１：データセットを作る

U2Netのデータセットは、
**・画像
・マスク画像（セグメンテーション対象物が白（255）、それ以外が黒（０）のピクセルになったオリジナル画像と同じサイズの画像）**
で構成されます。

画像とマスク画像は、以下のようにそれぞれのディレクトリに、同名で保存します。


```
my_dataset
   |
   |___images
   |      |__image0.jpg
   |      |__image1.jpg
   |      |__image2.jpg
   |
   |
   |___masks
   　　      |__image0.png
   　　      |__image1.png
   　　      |__image2.png
```

### 画像を集める
オリジナル画像を集めます。
フェイスマスクのようなそれほどパターンのないオブジェクトなら、２００枚程度でも十分だと思います。

### アノテーションツールでマスク画像を作る

アノテーションツールはなんでも構いませんが、今回は無料ツールの**labelMe**を使います。

https://github.com/wkentaro/labelme

以下のリンクから、各OS用のGUIアプリケーションを直接ダウンロードできます。

https://github.com/wkentaro/labelme/releases/tag/v5.0.5

labelMeのアプリケーションを開き、集めた画像のディレクトリを開きます。

![スクリーンショット 2022-11-15 0.42.04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/740f8381-d792-0c4d-814f-dc5235a43540.png)

CreatePolygonsボタンを押してから、対象オブジェクトの輪郭の頂点をクリックしていき、囲みます。
開始点まで来ると、ラベル名を選ぶウインドウが開くのでラベル名を選びます。

![スクリーンショット 2022-11-15 0.43.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/f9f4e03c-5c2b-f558-59cd-b7cda0d40c0c.png)

Saveボタンを押して、画像と同名のアノテーションjsonファイルを保存します。
image0.jpgなら、image0.jsonというファイルが保存されます。

次の関数でjsonファイルからマスク画像を作ります。

```python
import numpy as np
import math
from PIL import ImageDraw, Image

def shape_to_mask(
    img_shape, points, shape_type=None, line_width=10, point_size=5
):
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = Image.fromarray(mask)
    draw = ImageDraw.Draw(mask)
    xy = [tuple(point) for point in points]
    if shape_type == "circle":
        assert len(xy) == 2, "Shape of shape_type=circle must have 2 points"
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
    elif shape_type == "rectangle":
        assert len(xy) == 2, "Shape of shape_type=rectangle must have 2 points"
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == "line":
        assert len(xy) == 2, "Shape of shape_type=line must have 2 points"
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "linestrip":
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "point":
        assert len(xy) == 1, "Shape of shape_type=point must have 1 points"
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        assert len(xy) > 2, "Polygon must have points more than 2"
        draw.polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask
```

```python
import json
import cv2

dest_dir = 'my_dataset/masks/'

with open(json_path, "r",encoding="utf-8") as f:
  dj = json.load(f)
  # dj['shapes'][0]は今回一つのラベルのため。
  mask = shape_to_mask((dj['imageHeight'],dj['imageWidth']), dj['shapes'][0]['points'], shape_type=None,line_width=1, point_size=1)
  mask_img = mask.astype(np.int)#booleanを0,1に変換
  dest = os.path.join(dest_dir, os.path.splitext(os.path.basename(json_path))[0]+".jpg")
  cv2.imwrite(dest,mask_img*255)
```

![sorasanPAR55131_TP_V4 (1).jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/40e7b78d-1aad-abab-4132-f134a6146d49.jpeg)

# トレーニング

U2netのリポジトリをクローンします。

https://github.com/xuebinqin/U-2-Net

u2net_train.pyの以下の部分のパスを独自のものに書き換えます。
51行目のデータの親ディレクトリ（data_dir）
52行目の画像ディレクトリ（tra_image_dir）
53行目のマスクディレクトリ（tra_label_dir）
58行目のトレーニングしたモデルウェイト保存先

```python:u2net_train.py
data_dir = "my_data/"
tra_image_dir = 'images/'
tra_label_dir = 'masks/'
...
model_dir = 'my_saved_models/'
```

dataloader.pyの221行目を以下のように書き換えます。

```python:dataloader.py
##		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)}

## ↓ .copy() をつける

		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg.copy()), 'label': torch.from_numpy(tmpLbl.copy())}
```

トレーニングを開始します。

```shell
python u2net_train.py
```

u2net_train.py１１６行目のsave_frq（デフォルトは2000）のイテレーションごとに、モデルウェイトが保存されます。
保存したウェイトは通常の推論と同じ手順で使えます。

https://qiita.com/john-rocky/items/9bc43884f6c34ef47aff

# 途中で止めた学習を再開したい場合


```python:u2net_train.py
net.load_state_dict(torch.load(model_path))
```


python:u2net_train.pyの102行目に以下を追加してtrainを実行すると、途中のウェイトから学習を再開できます。
![nTmY8IS7byvleQu1668443720_1668443819.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/a79b2ba3-b093-014b-3023-5692fe509c05.png)



🐣
****
フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

機械学習、ARアプリ（Web/iOS）を作っています。
機械学習／AR関連の情報を発信しています。

[Twitter](https://twitter.com/JackdeS11)
[Medium](https://medium.com/@rockyshikoku)
[GitHub](https://github.com/john-rocky)