OpenCVで効率よくトリミングをしてみた

7917 ワード
Python OpenCV annotation Python テキストリンク

機械学習をするときにはまず教師データの前処理が大事だと言われています。
そして画像処理で物体認識をするときには学習したい物体だけトリミングして無駄を無くしたいですよね？
でもわざわざ画像編集ソフトを使ってトリミングして保存するのも手間です。
というわけで今回はOpenCVを使って効率よく手作業でトリミングをする方法を紹介します。

今回やること

ディレクトリ内の画像を一括で読み込んで、以下のように一枚ずつ学習したい領域をトリミングできるようにします。
トリミングした画像は保存用ディレクトリに格納されます。

切り取る領域の選択

まずはドラック＆ドロップで画像上の選択領域の原点になるx軸、y軸、高さ、幅の情報を取るにはcv2.selectROIを使うと領域選択する画面が簡単に立ちがります。
切り出す領域を選択したらEnterキーを押すと選択領域の座標情報をタプルで取得できます。
(Cキーを押すとキャンセルされて、(0. 0. 0, 0)が返されます)

selected = cv2.selectROI(img)

選択領域の座標情報を取得したら画像データを切り出すことができます。

imCrop = img[int(selected[1]):int(selected[1]+selected[3]),
                         int(selected[0]):int(selected[0]+selected[2])]

あとは切り出した画像データを保存するだけです。

cv2.imwrite(file_dir, imCrop)

完成したコード

ファイルのパス操作に関するコードも合わせた完成形はこちらです。
WindowsとUNIX系でも使えるようにパスの書式を統一しています。

trim_image.py

import cv2
import os
import sys

if __name__ == '__main__':
    args = sys.argv
    data_dir = args[1].replace('\\', '/')
    annotated_dir = data_dir + 'trimed'
    os.mkdir(annotated_dir)
    files = os.listdir(data_dir)
    img_files = [
        f for f in files if '.jpeg' in f or '.jpg' in f or '.png' in f]
    print(img_files)
    for img_file in img_files:
        # Read image
        img_dir = data_dir + img_file
        img = cv2.imread(img_dir)

        # Select ROI
        selected = cv2.selectROI(img)
        if sum(selected):
            # Crop image
            imCrop = img[int(selected[1]):int(selected[1]+selected[3]),
                         int(selected[0]):int(selected[0]+selected[2])]
            # write annotated img
            file_dir = annotated_dir + '/' + img_file
            cv2.imwrite(file_dir, imCrop)
            print('saved!')

以下のコマンドで画像データが保存されたディレクトリのパスを指定してディレクトリ内の画像を一括でトリミングできるようになります。
トリミングした画像データはディレクトリ内に作成したtrimedディレクトリ内に保存されます。

python trim_image.py path/to/img_dir