YOLOXでビデオ映像から人物の映っているフレームだけ抽出してみた

# ０．はじめに
ビデオ映像から人が映っているフレームだけを取り出したい、と思ったことはないでしょうか。１時間のビデオ映像で人が映っているのが３分の場合、そこを探して取り出す、というのは人手ではなかなか手間です。

こうした手間を削減するため、物体検出DNNを使って人物が映っているフレームのみを切り出したビデオ映像を再構成することで、映像チェックの負担を下げることを試みました。

物体検出DNNとしては、2021年7月に発表された[YOLOX]を使用しました。さまざまな大きさの事前学習モデルが公開されているので、実行環境や処理内容に応じて選択でき、また使い勝手がよかったためです。コードも[YOLOX github]上で[Apache License 2.0](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/LICENSE)で公開されているものを使いました。

ここでは、以下の流れで説明します。

- 作業の方針
- 人物検出コードの説明
- 実行
- 誤検出への対応
- まとめ

# １．作業の方針
この作業では、極力自前でつくらず、既存で使えるものを利用することとしました。従って、コードは[YOLOX github]で公開されている[tools/demo.py]をベースにし、[YOLOX github]で公開されている事前学習モデルを活用します。

[YOLOX github]の[Benchmarkセクション](https://github.com/Megvii-BaseDetection/YOLOX#benchmark)には[COCOクラス](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/data/datasets/coco_classes.py)に使える学習済みモデルが公開されています。パラメータ数1M弱の[yolox-nano](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_nano.py)から100M近いパラメータを持つ[yolox-x](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_x.py)までのラインアップがあるので、精度要件や使用できる実行環境のリソースに応じて適切なものが選択できます。

|モデル名		|パラメータ数（M）|学習済みモデル／weights	|
|:--:|--:|:--:|
|[yolox-nano]	|	0.91	|	[yolox_nano.pth]		|
|[yolox-tiny]	|	5.06	|	[yolox_tiny.pth]		|
|[yolox-s]		|	9.0		|	[yolox_s.pth]			|
|[yolox-m]		|	25.3	|	[yolox_m.pth]			|
|[yolox-l]		|	54.2	|	[yolox_l.pth]			|
|[yolox-x]		|	99.1	|	[yolox_x.pth]			|


[YOLOX github]ではカスタムトレーニングの方法についても紹介されており、比較的簡単に実行できます。

# ２．人物検出コードの説明
[YOLOX github]で提供されている[tools/demo.py]を`demo2.py`としてコピーし、以下の修正を施します。

### 2.1 パッケージインポート
YOLOXの推論結果を処理するために使うnumpyのインポートを追加します。
~~~python
import numpy as np
~~~

### 2.2 make_parser()関数
[tools/demo.py]の主な引数の説明は以下です。なお、[YOLOX]では、EXPERIMENTファイルで各モデルのコンフィグを定義します。提供されている学習済みモデルのコンフィグは、[YOLOX/exps/default/]配下に格納されています。

|引数|引数の説明|
|:--:|:--|
|demo|demoのタイプを `image`, `video`, `webcam` から指定。|
|-h, --help|helpメッセージ。|
|-n NAME, --name NAME|上記表の「モデル名」から指定すると`exps/default`配下から適切なEXPERIMENTファイルを自動で選択。|
|-f EXP_FILE, --exp_file EXP_FILE|-n, --nameのかわりにEXPERIMENTファイルを直接指定。カスタムで学習したモデルを使う場合に使う。|
|--path PATH|入力とする画像ファイル、あるいはビデオファイルへのパス。|
|--camid CAMID|webcamを使ったデモの場合のcamera id。|
|--save_result|画像/ビデオに対する推論結果を保存する場合に指定。保存先は出力メッセージの`Saving detection result in`の後に記載。|
|-c CKPT, --ckpt CKPT|使用する学習済みモデル。|
|--device DEVICE|モデルを実行するデバイスとして、`cpu` あるいは `gpu` を指定。省略時は`cpu`で実行。|

以下の２行を追加して、抽出したいクラスとビデオのフレームレートを指定できるようにします。
ここでビデオのフレームレートを指定できるようにしたのは、ビデオ開始からの経過時間をfpsを使って算出したかったのですが、使用したビデオ映像は5fpsにも関らず、`cv2.VideoCapture(path).get(cv2.CAP_PROP_FPS)`では5fpsと取得できなかったためです。`cv2.VideoCapture(path).get(cv2.CAP_PROP_FPS)`で本来のfpsが取得できる場合は、この引数は不要です。
~~~python
    parser.add_argument("--extract", type=str, default=None, help="extract a class")
    parser.add_argument("--frame_rate", type=int, default=0, help="frame rate")
~~~
|追加した引数|追加した引数の説明|
|:--:|:--|
|--extract|抽出したいクラスをCOCOクラスでの名前を使って指定します。人物を抽出する場合は、personを指定します。|
|--frame_rate|ビデオのフレームレートを指定します。指定しない場合は、cv2.VideoCapture(path).get(cv2.CAP_PROP_FPS)で取得したフレームレートを使います。|

### 2.3 Predictor()クラス
指定されたEXPERIMENTファイルに基づいた[YOLOX]モデルクラスの定義です。以下の２つのメソッドが使えます。

- inference()
  - 推論を実行します。
  - [tools/demo.py]では推論実行ごとにログを出力しますが、ビデオでは出力が大量になるので、コメントアウトします。
~~~python
            #logger.info("Infer time: {:.4f}s".format(time.time() - t0))
~~~


- visual()
  - inference()で検出したオブジェクトをバウンディングボックスで囲み、クラス名、スコアを表記します。

### 2.4 image_demo()関数
画像処理を行います。画像に対してもクラスを抽出するため、以下の変更を行います。
- image_demo()の第五引数をargsに変更します。
~~~python
def image_demo(predictor, vis_folder, path, current_time, args):
~~~
- predictor.inference()実行直後に以下を追加します。
~~~python
        outputs[0] = extract_cls(outputs[0], args.extract)
~~~
- save_resultをargsから取得するように変更します。
~~~python
        if args.save_result:
~~~

### 2.5 imageflow_demo()関数
ビデオ処理を行います。`while True`のループに入る前に、以下の準備をします。
~~~python
    frame_rate = int(fps) if args.frame_rate == 0 else args.frame_rate
    frame_num = -1
~~~
ループの中身を以下に変更します。`--extract`によってクラスの抽出が指定されている場合は、そのフレームがビデオのどの位置かを示すために`cv2.putText()`によって、ビデオの最初からの経過時間とフレーム番号を左上に表示するようにしています。
~~~python
    while True:
        ret_val, frame = cap.read()
        frame_num += 1
        if ret_val:
            outputs, img_info = predictor.inference(frame)
            outputs[0] = extract_cls(outputs[0], args.extract)
            result_frame = predictor.visual(outputs[0], img_info, predictor.confthre)

            if args.save_result:
                if args.extract is None:
                    vid_writer.write(result_frame)
                elif (outputs[0] is not None) and (args.extract is not None):
                    minutes = frame_num//(frame_rate*60)
                    seconds = (frame_num - minutes*(frame_rate*60)) // frame_rate
                    frame = cv2.putText(result_frame, "Time: {:d} min {:d} sec".format(minutes, seconds),
                                (60, 60), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (0, 255, 127), 2)
                    frame = cv2.putText(result_frame, "Frame Num: {:3d}".format(frame_num),
                                (60, 120), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (0, 255, 127), 2)
                    vid_writer.write(result_frame)

            else:
                cv2.namedWindow("yolox", cv2.WINDOW_NORMAL)
                cv2.imshow("yolox", result_frame)
            ch = cv2.waitKey(1)
            if ch == 27 or ch == ord("q") or ch == ord("Q"):
                break
        else:
            break
~~~

### 2.6 extract_cls()関数
新規関数です。`predictor.inference()`の出力から`--extract`で指定したクラスを抽出するよう、`image_demo()`、`imageflow_demo()`、それぞれから上述のとおりに呼び出します。
~~~python
def extract_cls(output, extract_class):
    if extract_class is not None:
        if output is not None:
            # クラス名はinfres[6]で確認
            extracted = [infres.cpu().detach().numpy() for infres in output
                            if infres[6] == COCO_CLASSES.index(extract_class)]
            if extracted != []:
                return torch.tensor(np.array(extracted)).cuda()
            else:
                return None
    else:
        return output
~~~

### 2.7 main()関数
EXPERIMENTファイルに基づいた[YOLOX]クラスのインスタンス化を行い、image、videoに応じてそれぞれ`image_demo()`、`imageflow_demo()`を呼び出します。
上記で`image_demo()`の定義で第五引数を`args`に変更したため、呼出し箇所も変更します。
~~~python
        image_demo(predictor, vis_folder, args.path, current_time, args)
~~~

# ３．実行

### 3.1 インストール
以下のようにして[YOLOX github]をクローンし、環境を準備します。
```
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -v -e .  # or  python3 setup.py develop
```

また、先ほど作成した`demo2.py`をローカルの`YOLOX`直下に配置します。

### 3.2 学習済みモデル

上記表のURLを使うことで、例えば、[yolox-l](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_l.py)の学習済みモデルは以下のように取得することもできます。

```
wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.pth
```


### 3.3 画像での物体検出
`IMAGES/nightroad.jpg`に対して物体検出を行います。以下の実行により、`Saving detection result in`の後ろに示されるパスに出力結果が格納されます。
```
python demo2.py image -n yolox-l -c yolox_l.pth --path IMAGES/nightroad.jpg --save_result --device gpu
```
以下のとおり、人、車、バスなどを検出しています。
なお、バウンディングボックスの色は`_COLORS` ([YOLOX/yolox/utils/visualize.py])でクラスごとに定義されていますが、以下の画像で人が目立ちやすいように、`person`クラスに対応するRGBを`0.000, 0.447, 0.741`から`0.000, 1.000, 0.500`に変更しています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/c6f96349-11bf-2c8e-f93e-6945a38bbf7a.png)

今度は、personを抽出するために`--extract person`を追加します。
```
python demo2.py image -n yolox-l -c yolox_l.pth --path IMAGES/nightroad.jpg --save_result --extract person --device gpu
```
以下のとおり、人のみをバウンディングボックスでマークしています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/fb9455e3-a19c-820c-0fc6-4f9bad287dee.png)

### 3.4 ビデオでの物体検出

さあ、いよいよビデオに対して実行をかけます。例えば以下のように[yolox-l]を使い実行すると、人を検出したフレームのみ`video save_path is`で示されるファイルにセーブされます。

```
python demo2.py video -n yolox-l -c yolox_l.pth --path video.avi --save_result --extract person --frame_rate 5 --device gpu
```
```
2022-04-09 17:52:50.241 | INFO     | __main__:main:290 - Args: Namespace(camid=0, ckpt='yolox_l.pth', conf=0.3, demo='video', device='gpu', exp_file=None, experiment_name='yolox_l', extract='person', fp16=False, frame_rate=5, fuse=False, legacy=False, name='yolox-l', nms=0.3, path='video.avi', save_result=True, trt=False, tsize=None)
2022-04-09 17:52:50.944 | INFO     | __main__:main:300 - Model Summary: Params: 54.21M, Gflops: 155.65
2022-04-09 17:52:53.918 | INFO     | __main__:main:313 - loading checkpoint
2022-04-09 17:52:54.293 | INFO     | __main__:main:317 - loaded checkpoint done.
2022-04-09 17:52:54.406 | INFO     | __main__:imageflow_demo:227 - video save_path is ./YOLOX_outputs/yolox_l/vis_res/2022_04_09_17_52_54/video.avi

```

# ４．誤検出への対応

ビデオ映像には、さまざまなものが映り込んだり、夜間の照明でものの見え方が昼間と違う場合があります。今回の方針として、学習済みモデルを使うことにしたので、残念ながら誤検出が起こる可能性があります。

市販のAI製品では、誤検出を防ぐために膨大なデータセットを使って学習しているようですが、ここでは、今回の方針の範囲での誤検出の対応案をいくつかあげたいと思います。

### 4.1 大きな学習済みモデルを試す
[YOLOX github]ではいくつかの学習済みモデルを開示していることは先に述べました。やはり大きなモデルは精度が良いので、false negativeの対策として試してみるのが良いと思います。ただし、大きな学習済みモデルは計算時間がより長くなりますし、経験的にfalse positiveには必ずしも有効ではないようです。

### 4.2 閾値を変更する
検出の閾値を調整することで誤検出を減らすことができる場合があります。以下が閾値として設定変更できます。
|閾値設定の引数	|引数の説明	|
|:--:			|:--		|
|--conf CONF	|CONF未満のスコアのバウンディングボックスを捨てる。デフォルトは0.3。	|
|--nms NMS		|Non Maximum Suppression処理はmegengine.functional.nn.nmsを使用。NMSで指定する値は、オーバーラップする部分のIoUの閾値。デフォルトは0.3。	|

### 4.3 マスクをかける
誤検出する箇所が固定の場合は、その箇所にマスクをかけてYOLOXに入力するのも一案です。ただし、そこに人が映っている場合は検出できなくなりますし、条件によって誤検出箇所が移動する場合もあります。


# まとめ
事前学習モデルや公開されているコードを使うことで、比較的容易にビデオ映像チェックの負担を下げることができました。
このように、公開されているデータやコードだけで身の回りの困りごとにAIが適用できるのは何ともうれしいことなぁ、と思いました。



[YOLOX]: https://arxiv.org/abs/2107.08430
[Yolov4]: https://arxiv.org/abs/2004.10934
[YOLOX github]: https://github.com/Megvii-BaseDetection/YOLOX
[COCOクラス]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/data/datasets/coco_classes.py


[yolox-nano]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_nano.py
[yolox_nano.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.pth
[yolox-tiny]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_tiny.py
[yolox_tiny.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny.pth
[yolox-s]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_s.py
[yolox_s.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth
[yolox-m]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_m.py
[yolox_m.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m.pth
[yolox-l]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_l.py
[yolox_l.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.pth
[yolox-x]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_x.py
[yolox_x.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth
[yolov3]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolov3.py
[yolox_darknet.pth]: https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_darknet.pth

[tools/demo.py]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/tools/demo.py
[YOLOX/exps/default/]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default
[YOLOX/yolox/utils/visualize.py]: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/utils/visualize.py