【YOLOv5】マスクの物体検出

# はじめに
YOLOv5を用いての物体検知について勉強したので、マスクを着用しているかどうかを検出するモデルを作成してみました。

# ライブラリの準備

今回は、pytorchで書かれたyolo v5を用いるので、以下のコマンドでgitからクローンしてきます。
pythonのバージョンは`3.8`以上だと良いみたいです。

```
$ git clone https://github.com/ultralytics/yolov5.git
```
次に、必要なライブラリをインストールします。

```
$ cd yolo5
$ pip install -r requirements.txt
```

# データの準備

アノテーション済みのマスク画像データを提供している[Mask Wearing Dataset](https://public.roboflow.com/object-detection/mask-wearing)からデータをダウンロードします。
ダウンロードボタンを押して、以下の画像のように、`YOLO Darknet`を選択します。
![スクリーンショット 2021-10-26 4.07.51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/6d8638d1-5dae-bdff-a208-8fc20b95254a.png)
ダウンロードしたデータは、`train`, `valid`, `test`の３つのフォルダに分かれているので、それらのフォルダを`yolo5/data/mask_data/`に移動させてください。

# 学習

`data.yaml`という、学習を行う時の訓練データと検証データのパスや、検出クラス数やクラス名についてを記載するファイルを用意し、`yolo5/data.yaml`に置きます。

```data.yaml
train: data/mask_data/train
val: data/mask_data/valid

nc: 2

names: ["mask", "no-mask"]
```

これで準備完了です。以下のコマンドで学習開始します。

```
$ python train.py --batch 16 --epochs 100 --data data.yaml
```
学習が終わると、モデルの重みが以下のファイルに保存されます。

```
Optimizer stripped from runs/train/exp/weights/last.pt, 14.5MB
Optimizer stripped from runs/train/exp/weights/best.pt, 14.5MB
```

# テスト

学習した重みを用いて、以下の２つの画像のマスク着用を検知してみます。
![try2.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/bede2e25-64df-0c0e-41bc-284f2ebd6fe2.jpeg)

![try.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/333a827f-4ae1-a695-2d9d-60d4706a5dec.jpeg)

テスト画像を`yolo5/example_data/`に置いて、以下のコマンドを実行します。
変数は、`--source テスト画像ディレクトリ`, `--weights 学習モデルの重み`, `--conf-thres Confidense scoreの閾値`を指定し、`--save-txt`を加えることで、検出したバウンディングボックスの位置と分類したクラスについて書かれたテキストファイルを作成してくれます。

```
$ python detect.py --source example_data/ --weights runs/train/exp/weights/best.pt --conf-thres 0.4 --save-txt
```
これによって得られた結果が以下の通りです。
![low2.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/f664e91c-7273-8d5b-43dd-0a3043b757b0.jpeg)
![low.jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/515d4296-c0e6-5736-a3fc-24768610324d.jpeg)
1枚目はうまく検出できていますが、2枚目は誤認識が目立ちます。
今回のデータセットは149枚と少ないため、あまり精度が出なかったように考えられます。
ということで、データの水増しを行いました。

# データの水増し（Data Augmentation）
前に自分が書いた記事である[ImageDataGenaratorを用いた画像の水増し](https://qiita.com/yassyyuki/items/f1ecddea04ca728c765c)を使用しても良かったのですが、今回は他の方法も試してみるため、`alubumentations`というライブラリを使用してみることにしました。
`albumentations`についての使い方は以下の方の記事が視覚的でとてもわかりやすかったのでおすすめです。
[Albumentationsのaugmentationをひたすら動かす](https://qiita.com/kurilab/items/b69e1be8d0224ae139ad)

以下のコマンドでライブラリのインストールをします。

```
$ pip install albumentations
```
以下のような`augmentation.py`を作成しました。
アノテーション作業が面倒なため、今回は楽をするために回転や歪み、移動などの処理での水増しを行わずに、`RGB Shift`や`Blur`といった、座標自体の変換がない処理を加えて、アノテーションファイルを使い回しました。

```augmentation.py
import albumentations as albu
import glob
import cv2
import shutil
from PIL import Image


def get_augmentation_1():
    train_transform = [
        albu.Blur(p=1)
    ]
    return albu.Compose(train_transform)

def get_augmentation_2():
    train_transform = [
        albu.RGBShift(p=1)
    ]
    return albu.Compose(train_transform)

def get_augmentation_3():
    train_transform = [
        albu.HueSaturationValue(p=1)
    ]
    return albu.Compose(train_transform)


IMG_DIR = './mask_data/train/'
imgs = glob.glob(IMG_DIR + '*.jpg')
txts = glob.glob(IMG_DIR + '*.txt')

# ソートをすることで画像とテキストファイルの対応している順番を揃える。
imgs.sort()
txts.sort()

num = 0
for i in range(len(imgs)):
    img_origin = cv2.imread(imgs[i])
    img_origin = cv2.cvtColor(img_origin, cv2.COLOR_BGR2RGB)
    img = img_origin.copy()
    transforms_1 = get_augmentation_1()
    transforms_2= get_augmentation_2()
    transforms_3 = get_augmentation_3()

    augmented_1 = transforms_1(image=img)
    augmented_2 = transforms_2(image=img)
    augmented_3 = transforms_3(image=img)

    img_1 = augmented_1['image']
    img_2 = augmented_2['image']
    img_3 = augmented_3['image']
    
    img_pil_1 = Image.fromarray(img_1)
    img_pil_2 = Image.fromarray(img_2)
    img_pil_3 = Image.fromarray(img_3)

    img_pil_1.save(f"./data/mask_data/train/blur_{num}.jpg")
    img_pil_2.save(f"./data/mask_data/train/rgb_shift_{num}.jpg")
    img_pil_3.save(f"./data/mask_data/train/hue_saturation_value_{num}.jpg")
    
    shutil.copy(txts[i], f"./data/mask_data/train/blur_{num}.txt")
    shutil.copy(txts[i], f"./data/mask_data/train/rgb_shift_{num}.txt")
    shutil.copy(txts[i], f"./data/mask_data/train/hue_saturation_value_{num}.txt")
    
    num += 1
```

実行し、画像の水増しを行います。

```
$ python augmentation.py
```
これで、データ数が4倍になったはずです。


# 学習（データ水増し後）
そして、学習のやり直しを行います。

```
$ python train.py --batch 16 --epochs 100 --data data.yaml
```

学習が終わると、重みが以下のファイルに作成されました。

```
Optimizer stripped from runs/train/exp2/weights/last.pt, 14.5MB
Optimizer stripped from runs/train/exp2/weights/best.pt, 14.5MB
```

# テスト
結果は以下の通りです。
![try2 (1).jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/a2297f1c-2ad9-dc90-4a72-98e6baea87eb.jpeg)
![try (1).jpeg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/298735/64389aca-b065-ab3d-da9c-5391198ae3e3.jpeg)
1枚目の画像はConfidense Scoreが上昇し、2枚目の画像は誤認識はあるものの、1回目と比べると検出漏れがなく、精度も向上しているのがみて取れます。
データを増やし、エポック数などのハイパーパラメータを調節すればもっと良い精度が出ることが期待される結果となりました。

# おわりに
今回はYOLOv5の学習と検証、データの水増しについて書かせていただきました。
こんなに簡単なのに精度の高い結果を出すYOLOv5はすごいですね。。
物体検出はなんといってもアノテーションが面倒な作業かと思われます。
今回のように、現在アノテーション済みのデータを、座標変換がないように処理することでアノテーションの作業を省略することはある程度効果的です。
しかし、理想はもっと色々な手法でデータの水増しを行って、`LabelImg`などのツールを用いてアノテーションを行い、頑健なモデルになるようにするべきでしょう。
最後までお読みいただき、ありがとうございました。