Python + OpenCVで野球ボールをトラッキング

久しぶりにOpenCVのサイトを見たら新しくなってたので、ちょっと触ってみました。

だいぶ前にJavaで野球ボールをトラッキングをしていたので、今回はPythonでやってみました。
トラッキングは、背景差分とテンプレートマッチングを組み合わせて行います。

環境構築

環境は以下になります。

• Python : 3.8.0
• OpentCV : 4.2.0

venvで環境を作っています。venvじゃなくても問題ないです。

$ python -m venv pythonOpevCV
$ cd pythonOpenCV
$ Script\activate
$ python -m pip install --upgrade pip
$ pip install opencv-python
Successfully installed numpy-1.18.2 opencv-python-4.2.0.34

numpyもインストールされるんや。

用意するもの

以下のものを用意してください。

• 野球ボールが映る動画
• 野球ボールを切り取った画像（テンプレート画像）

ボールを切り取った画像は以下のようなもので大丈夫です。以下見本です。

手順

以下の手順でトラッキングを行っていきます。

1. 動画のフレームごとに画像で保存
2. フレーム画像とテンプレート画像をグレースケール化
3. フレーム画像とテンプレート画像を2値化
4. 前後のフレーム画像で背景差分を行う
5. 背景差分の結果の画像からテンプレートマッチングを行う
6. トラッキングを結果を描画

いきなりオリジナルのフレーム画像に対してテンプレートマッチングを行っても、雲や背景の建物をボールと間違えて検出してしまい、精度が良くないのでグレースケール→2値化で白黒の画像にします。その白黒画像の状態で前後のフレームで背景差分を行うと、雲や背景は0.1秒ぐらいではほとんど動かないので移動しているボールがはっきり検出できます。

グレースケールした画像が以下のようになります。

赤枠で囲ってある範囲にあるボールを検出します。この画像を2値化し白黒の画像にします。

ボールを白色で検出できていますが、背景とボールが同じ白色のためボールだけを検出できません。ここで前後のフレーム画像で背景差分を行うと以下のようになります。

背景の白色は動作していないので背景差分では検出されず、ボールとその他のノイズのみが白色で検出されています。この状態で以下のようなテンプレート画像を2値化した画像とテンプレートマッチングを行います。

こうすることで光度や背景の影響を受けずにボールを検出することができます。

ソースコード

ソースコードは以下のようになります。VIDEOPATHに対象の動画を、TEMPLATEPATHにテンプレート画像を配置してください。

main.py

import glob
import re
import cv2


VIDEOPATH = "media/video/video.mp4"
IMAGEPATH = "media/image/"
TEMPLATEPATH = "template.jpeg"


def save_frames(video_path, image_dir):
    """
    動画からフレームの画像を抽出
    """
    cap = cv2.VideoCapture(video_path)
    digit = len(str(int(cap.get(cv2.CAP_PROP_FRAME_COUNT))))
    n = 0
    while True:
        ret, frame = cap.read()
        if ret:
            cv2.imwrite("{}original/frame_{}.{}".format(IMAGEPATH, n, "jpeg"), frame)
            n += 1
        else:
            return


def do_grayscale(image_path):
    """
    画像をグレースケール化
    """
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    save_image(image_path, "gray", gray)


def do_binarization(image_path):
    """
    画像を2値化
    """
    img = cv2.imread(image_path)
    ret, img_thresh = cv2.threshold(img, 100, 255, cv2.THRESH_BINARY)
    save_image(image_path, "binary", img_thresh)


def do_backgroundsub():
    """
    背景差分を行う
    """
    img_list = glob.glob(IMAGEPATH + "binary/frame*.jpeg")
    num = lambda val: int(re.sub("\D","",val))
    sorted(img_list,key=(num))
    source = img_list[0]
    for path in img_list:
        diff = cv2.absdiff(cv2.imread(source),cv2.imread(path))
        source = path
        save_image(path, "bgsub", diff)


def do_template_matching():
    """
    テンプレート画像とフレーム画像でテンプレートマッチングを行う
    """
    template_img = cv2.imread(IMAGEPATH + "binary/" + TEMPLATEPATH)
    img_list = glob.glob(IMAGEPATH + "bgsub/frame*.jpeg")
    num = lambda val: int(re.sub("\D","",val))
    sorted(img_list,key=(num))
    location_list = []
    for path in img_list:
        result = cv2.matchTemplate(cv2.imread(path), template_img, cv2.TM_CCOEFF)
        minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(result)
        location_list.append(maxLoc)
    return location_list

def draw_rectangle(location_list):
    """
    マッチング結果を画像に描画する
    """
    source = cv2.imread(IMAGEPATH + "original/frame_0.jpeg")
    cv2.imwrite(IMAGEPATH + "result.jpeg",source)
    source = cv2.imread(IMAGEPATH + "result.jpeg")
    for loc in location_list:
        lx, ly, rx, ry = loc[0] - 10, loc[1] - 10, loc[0] + 10, loc[1] + 10
        img = cv2.rectangle(source, (lx, ly), (rx, ry), (0, 255, 0), 3)
        cv2.imwrite(IMAGEPATH + "result.jpeg",img)

def save_image(img_path, dir, img):
    """
    画像を保存する
    img_path : 画像のパス
    dir : ディレクトリ名
    img : 画像データ
    """
    file_name = img_path.replace("\\","/").split(".")[0].split("/")[-1]
    cv2.imwrite("{}{}/{}.{}".format(IMAGEPATH, dir, file_name,"jpeg"), img)


if __name__=="__main__":
    # ①動画をフレームごとに分割
    save_frames(VIDEOPATH,IMAGEPATH)
    # ②テンプレート画像とフレーム画像をグレースケール化
    do_grayscale(IMAGEPATH + TEMPLATEPATH)
    for path in glob.glob(IMAGEPATH + "original/*.jpeg"):
        do_grayscale(path)
    # ③テンプレート画像とフレーム画像の2値化
    for path in glob.glob(IMAGEPATH + "gray/*.jpeg"):
        do_binarization(path)
    # ④背景差分を行う
    do_backgroundsub()
    # ⑤テンプレートマッチングを行う
    location_list = do_template_matching()
    # ⑥マッチングした座標を投影
    draw_rectangle(location_list)

結果

投球動画でボールを検出してみました。概ね検出できていますが、ボールの軌道以外の部分も検出してます。

ノイズや外れ値を修正する方法があったのですが忘れたのでまた思い出します。

まとめ

Python + OpenCVで背景差分とテンプレートマッチングを用いて野球ボールのトラッキングを行いました。
YOLOを使えば、画像の中から野球ボールを検出することができるようなので、このあたりも試してみたいです。