【OpenCV】画像からマウスでバウンディングボックスを選択して保存する

概要

最近機械学習モデルの学習用に画像データを収集しているのですが，
その際，画像に写っている物体のみを切り取って保存したい状況がありました．
今回はこちらをOpenCVを使用して実装してみたので，メモを残しておきます．
今回実装したコードはこちらに公開しております．

cv2.selectROI

OpenCVでは，画像からインタラクティブに矩形を選択できる関数として，cv2.selectROIが提供されています．(参考)
今回はこちらを使用して，矩形を選択し，選択した矩形内の画像を保存するようにしました．

以下はcv2.selectROIを使用して，選択した矩形内を別のwindowに表示するサンプルコードです．

import cv2
import numpy as np

img = cv2.imread("cat.jpg")

bbox = cv2.selectROI(img, fromCenter=False, showCrosshair=False)

# crop image
left, top, width, height = map(int, bbox)
cropped_img = img[top : top + height, left : left + width]

cv2.imshow("cropped image", cropped_img)
cv2.waitKey(0)

上記サンプルを実際に動かした結果

今回実装したコード

今回実装するコードの内容は以下の通りです．

1. 画像を読み込む
2. 画像からバウンディングボックスの範囲をマウスで指定し，Enterでbbox内の画像のみのプレビューを表示
3. 選択したバウンディングボックスを保存する場合はsキーで保存，囲み直す場合はrキーを押す
4. Ctrl + cかEscapeが押されるまで，2-3を繰り返す

import argparse
import glob
import os

import cv2


def get_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("img_path", type=str)
    parser.add_argument("--save_dir", type=str, default="./result")
    return parser.parse_args()


def crop_bbox(image_path: str, save_dir: str) -> None:
    img = cv2.imread(image_path)
    name = os.path.splitext(os.path.basename(image_path))[0]

    # count the number of saved images from the image
    cnt = len(glob.glob(os.path.join(save_dir, f"{name}*.jpg")))

    while True:
        # select region of interest (bounding box)
        ROI = cv2.selectROI(
            "Please draw a bounding box", img, fromCenter=False, showCrosshair=False
        )

        left, top, width, height = map(int, ROI)

        # check if a bbox is valid or not
        if width == 0 or height == 0:
            break

        # cropping image
        cropped_img = img[top : top + height, left : left + width]
        cv2.imshow("crop", cropped_img)

        while True:
            k = cv2.waitKey(0)
            if k == ord("r"):
                # retry selecting bboxes
                cv2.destroyWindow("crop")
                break
            elif k == ord("s"):
                # save cropped images
                cnt += 1
                save_path = os.path.join(save_dir, f"{name}{cnt:0>3}.jpg")
                cv2.imwrite(save_path, cropped_img)
                cv2.destroyWindow("crop")
                break


def main():
    args = get_arguments()

    os.makedirs(args.save_dir, exist_ok=True)

    if not os.path.exist(args.img_path):
        raise FileNotFoundError
    else:
        crop_bbox(args.img_path, args.save_dir)


if __name__ == "__main__":
    main()

動作確認

画像中の選択した範囲を，新たな画像として保存できていることが確認できました！

参考

• OpenCV: High-level GUI
• How to select a bounding box ( ROI ) in OpenCV (C++/Python) ? | LearnOpenCV #
• yiskw713/cv_utils: computer visoin utilities