魚眼レンズで角度のある人物を検出する機械学習モデル「RAPiD」を動画対応してみた。

2021.07.27

せーのでございます。

複数の人物を検出、追跡する機械学習モデルというのは結構見つけられるのですが、それがtop-view、つまり天井からのカメラとなると急にその数は減っていきます。
カメラの角度は斜めで人物の全身が確認できるもの、街の監視カメラくらいだと人物検出しやすいのですが、例えばコンビニのカメラや大きめの図書館の天井カメラのようなものになるとまずそれが人だと検出するのが難しく、ましてや速度が求められるリアルタイム検出になるとかなり難しくなります。

今回はそんな天井からのカメラ、さらに人物が歪んで回転する魚眼レンズでも人物を検出できる「RAPiD」を触ってみました。

RAPiDの特徴

RAPiDは「Rotation-Aware People Detection in overhead fisheye images」の略です(なんか順番おかしい気もしますが)。つまり「角度を意識した」人物検出、ということです。
こちらはYOLOをベースにCOCOデータセットで学習をおこなった後に魚眼レンズにて追加学習を行い、その際に画像の中心とそこからの距離、角度の情報を加えています。そうすることで自動的にその人物をバウンディングボックスで囲った時の角度まで推論できるようになりました。その角度を推論するのも含めた損失関数を拡張(周期的損失関数)しているため、推論は一回で済み、YOLOと推論する速さはほぼ変わらないというものです。すごいですね。

なのですが、今回は角度情報より「俯瞰からのカメラで人物を高速に検出」という点が良い、と思い、使ってみることにしました。

静止画 => 動画へ

では早速やってみます。OSですがAnacondaを使っているので特に問いません。ただCUDAを使っているのでNVIDIA GPUが載ってるマシンが良さそうです。ちなみに私は最初WindowsのWSL2にて動かしていたのですが、Ubuntuに変えてみたところ処理速度が一気に上がったのでUbuntuを使っています。

該当ソースをGithubからダウンロードし、ReadmeにあるままにAnacondaを使って環境を構築します。CUDA10.2を普段使っていたのでそれを使ってみたところPyTorchの互換性が合わずにエラーが起きたので、素直にReadme通りCUDA10.1を使います。

conda create --name RAPiD_env python=3.7
conda activate RAPiD_env

conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
conda install -c conda-forge pycocotools
conda install tqdm opencv

# cd the_folder_to_install
git clone https://github.com/duanzhiihao/RAPiD.git

インストールが終わり、 事前に学習されたモデルをダウンロードしてexample.pyを叩いてみました。

うん、いい感じですね。ではこれを動画対応してみます。

もともとのソースは

from api import Detector

# Initialize detector
detector = Detector(model_name='rapid',
                    weights_path='./weights/pL1_MWHB1024_Mar11_4000.ckpt')

# A simple example to run on a single image and plt.imshow() it
detector.detect_one(img_path='./images/exhibition.jpg',
                    input_size=1024, conf_thres=0.3,
                    visualize=True)

このようになっています。このdetector.detect_oneのimg_pathを他の画像パスに変えればいろいろな画像を推論できるわけですね。
でもいちいちjpg画像を動画から出力して推論するのは効率が悪いです。他に方法はないでしょうか。

そうするとこのdetectorメソッドにもう一つインプット方法がありました。

    def detect_one(self, **kwargs):
        '''
        Inference on a single image.
        Args:
            img_path: str or pil_img: PIL.Image
            input_size: int, input resolution
            conf_thres: float, confidence threshold
            return_img: bool, if True, return am image with bbox visualizattion. \
                default: False
            visualize: bool, if True, plt.show the image with bbox visualization. \
                default: False
        '''
        assert 'img_path' in kwargs or 'pil_img' in kwargs
        img = kwargs.get('pil_img', None) or Image.open(kwargs['img_path'])
......

ご丁寧にコメントで書かれていました。つまり「img_path」という引数なら画像がおいてあるパスを指定するけれど「pil_img」という引数であればPIL、つまりpythonのpillowイメージをそのまま渡せるということです。

そこで動画からOpenCVで読み取ったフレームをPIL化して渡してみました。

from api import Detector
import numpy as np
import cv2

def cv2pil(imgCV):
    imgCV_RGB = cv2.cvtColor(imgCV, cv2.COLOR_BGR2RGB)
    imgPIL = Image.fromarray(imgCV_RGB)
    return imgPIL  

# Initialize detector
detector = Detector(model_name='rapid',
                    weights_path='./weights/pL1_MWHB1024_Mar11_4000.ckpt')

cap = cv2.VideoCapture('videos/test.mp4')

while(cap.isOpened()):
    # フレームを取得
    ret, frame = cap.read()

    if ret == False:
        break

    np_img = detector.detect_one(pil_img=cv2pil(frame),
                    input_size=1024, conf_thres=0.3,
                    return_img=True)
    new_image = np.array(np_img, dtype=np.uint8)
    np_img = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
    # フレームを表示
    cv2.imshow("Frame", np_img)
    
    # qキーが押されたら途中終了
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

そうするとこんなエラーが出てきました。

_predict_pil() got multiple values for argument 'pil_img'

pil_imgという変数が複数の場所から変更されている、とのこと。pil_imgを渡しているapi.pyを見てみると

    def _predict_pil(self, pil_img, **kwargs):
        '''
        Args:
            pil_img: PIL.Image.Image
            input_size: int, input resolution
            conf_thres: float, confidence threshold
        '''
        input_size = kwargs.get('input_size', self.input_size)
        conf_thres = kwargs.get('conf_thres', self.conf_thres)
        assert isinstance(pil_img, Image.Image), 'input must be a PIL.Image'
        assert input_size is not None, 'Please specify the input resolution'
        assert conf_thres is not None, 'Please specify the confidence threshold'

        # pad to square
        input_img, _, pad_info = utils.rect_to_square(pil_img, None, input_size, 0)

ここの部分で[pil_img]という変数名が[kwargs]の中にも入っているために起こっているようでした。なのでコードをさくっと

    def _predict_pil(self, pil_image, **kwargs):
        '''
        Args:
            pil_image: PIL.Image.Image
            input_size: int, input resolution
            conf_thres: float, confidence threshold
        '''
        input_size = kwargs.get('input_size', self.input_size)
        conf_thres = kwargs.get('conf_thres', self.conf_thres)
        assert isinstance(pil_image, Image.Image), 'input must be a PIL.Image'
        assert input_size is not None, 'Please specify the input resolution'
        assert conf_thres is not None, 'Please specify the confidence threshold'

        # pad to square
        input_img, _, pad_info = utils.rect_to_square(pil_image, None, input_size, 0)

と変更しました。そうすると

無事動きました!

[追記] エラーの件について開発者の方にissueを出したところ「シンプルにkwargsの中身を削除しちゃうね」ということでコミットしてくれましたので、現在はソースの変更なしに動きます。

まとめ

角度のついているバウンディングボックスというのは珍しいので色々と使えそうですね。
何より天井からのカメラに対して精度がいいです。魚眼レンズで学習させるとパフォーマンスが上がるようですが、通常画像だとどうなんでしょうか。トライしてみたいと思います。

参考リンク