YOLOラベルから牌画像を自動で切り出す｜Pythonで矩形座標をクロップする方法

YOLO

前回は、Tkinterを使ってYOLO用の矩形アノテーションツールを作成しました。

通常であれば、

画像を撮影
↓
矩形アノテーション
↓
YOLOで学習

と、そのまま学習へ進むことが多いと思います。

しかし今回は、ここから少し違う方法で学習データを作っていきます。

私が今回考えたデータ作成の流れでは、アノテーションした画像をそのまま学習に使うだけではなく、矩形アノテーションした麻雀牌を一度画像として切り出します。

そして、この切り出した牌画像を素材として使い、

撮影画像
↓
矩形アノテーション
↓
牌画像を切り出す
【今回】
↓
牌の柄部分を4点アノテーション
↓
別のベース画像へ射影変換して貼り付ける
↓
YOLOラベルを自動生成
↓
大量の学習画像を作る
↓
YOLOで学習

という流れでデータセットを作成します。

つまり今回のcrop.pyは、単純に画像を切り抜くためだけのプログラムではなく、この後の学習データ自動生成で使用する「牌素材」を作るためのツールです。

記事の後半には、今回使用しているcrop.pyのコード全文も掲載しています。

なぜ一度牌画像を切り出すのか

普通にYOLOを学習するだけなら、前回作成した矩形ラベルと撮影画像があれば、そのまま学習できます。

それでも一度牌画像を切り出すのは、少ない撮影画像から、より多くの学習画像を作りたいからです。

今回の方法では、撮影した麻雀牌を素材として取り出し、その後、別に用意したベース画像へ貼り付けて新しい画像を生成します。

例えば、同じ牌素材でも、

  • 貼り付ける位置
  • 牌の角度
  • ベースとなる画像
  • 画像処理やデータ拡張

などを変えることで、元画像とは異なる学習画像を作れます。

そのため、

撮影した画像をそのまま増やしていくのではなく、撮影した牌を「再利用できる素材」に変える

というのが今回の考え方です。

この方法が本当に最適かは試しながら改善していく必要がありますが、麻雀牌の学習データを効率よく増やすために、今回はこの独自の手順で進めています。

今回のプログラムでやること

crop.pyの処理自体はかなりシンプルです。

画像を読み込む
↓
同じ名前のYOLOラベルを読み込む
↓
YOLO座標を画像上の座標へ戻す
↓
矩形部分を切り出す
↓
牌画像として保存する

例えば、次のファイルがあるとします。

sample01.jpg
sample01.txt

sample01.txtに記録されている矩形を読み取り、元画像の対応する部分を自動でクロップします。

YOLO座標を画像上の座標へ戻す

YOLOのラベルは、そのままではPillowのcrop()には使用できません。

YOLOでは矩形が、

class_id x_center y_center width height

という形式で保存されています。

さらに座標はピクセルではなく、画像サイズに対して0〜1へ正規化されています。

そこで、まず元画像のサイズを使ってピクセル座標へ戻します。

def yolo_to_bbox(x_center, y_center, width, height, img_w, img_h):
    x_center *= img_w
    y_center *= img_h
    width *= img_w
    height *= img_h

    x1 = int(x_center - width / 2)
    y1 = int(y_center - height / 2)
    x2 = int(x_center + width / 2)
    y2 = int(y_center + height / 2)

    return x1, y1, x2, y2

例えば、

中心座標
＋
矩形の幅・高さ

として保存されているYOLOラベルを、

左上 (x1, y1)
右下 (x2, y2)

へ変換しています。

これでPillowから直接切り出せる座標になります。

矩形部分を牌画像として切り出す

座標へ変換したら、Pillowのcrop()を使って牌部分を切り出します。

crop = img.crop((x1, y1, x2, y2))

その前に、座標が画像外へ出ないように補正しています。

x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(img_w, x2)
y2 = min(img_h, y2)

アノテーションされた矩形が画像端に近い場合でも、安全にクロップできるようにしています。

切り出した画像は、

save_path = os.path.join(
    output_dir,
    f"{base_name}_{class_id}.jpg"
)

として保存します。

例えば、

sample01_0.jpg
sample01_8.jpg
sample01_34.jpg

のように、元画像の名前とクラスIDからファイル名を付けています。

これによって、あとからどの牌クラスなのか判別しやすくしています。

フォルダ内の画像をまとめて処理する

今回は1枚ずつ実行するのではなく、指定したフォルダ内のラベルをまとめて処理します。

for file in os.listdir(input_dir):
    if not file.endswith(".txt"):
        continue

.txtファイルを見つけたら、同じ名前の.jpgを探します。

txt_path = os.path.join(input_dir, file)
img_name = os.path.splitext(file)[0] + ".jpg"
img_path = os.path.join(input_dir, img_name)

例えば、

001.txt

があれば、

001.jpg

を読み込みます。

あとはラベルに含まれる矩形を順番にクロップしていくだけです。

そのため、一度アノテーションが終わっていれば、フォルダ内の牌画像をまとめて生成できます。

この切り出した画像をどう使うのか

今回作った牌画像は、これだけで完成ではありません。

次の工程では、切り出した牌画像から牌の柄部分を指定するための4点アノテーションを行います。

その4点を使って射影変換し、別に用意したベース画像へ貼り付けます。

最終的には、

1枚の牌素材
↓
さまざまな位置・形へ変換
↓
ベース画像へ貼り付け
↓
新しいYOLOラベルも自動生成
↓
多数の学習画像を生成

という仕組みにしていきます。

今回のcrop.pyは、その自動生成に使用する牌素材を準備する工程という位置付けです。

次回以降、この処理がデータセット自動生成へつながっていきます。

完成したコード

今回使用したcrop.pyの全文です。

コードが長いため、デフォルトでは折りたたんでいます。必要に応じて「crop.py の全文を見る」をクリックしてください。

crop.py の全文を見る 
import os
from PIL import Image

def yolo_to_bbox(x_center, y_center, width, height, img_w, img_h):
    x_center *= img_w
    y_center *= img_h
    width *= img_w
    height *= img_h

    x1 = int(x_center - width / 2)
    y1 = int(y_center - height / 2)
    x2 = int(x_center + width / 2)
    y2 = int(y_center + height / 2)

    return x1, y1, x2, y2


def crop_from_yolo(txt_path, img_path, output_dir):
    img = Image.open(img_path)
    img_w, img_h = img.size

    base_name = os.path.splitext(os.path.basename(img_path))[0]

    with open(txt_path, "r") as f:
        lines = f.readlines()

    for line in lines:
        parts = line.strip().split()
        if len(parts) != 5:
            continue

        class_id, x, y, w, h = parts
        class_id = int(class_id)
        x, y, w, h = map(float, (x, y, w, h))

        x1, y1, x2, y2 = yolo_to_bbox(x, y, w, h, img_w, img_h)

        # 範囲補正
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(img_w, x2)
        y2 = min(img_h, y2)

        crop = img.crop((x1, y1, x2, y2))

        save_path = os.path.join(output_dir, f"{base_name}_{class_id}.jpg")
        crop.save(save_path)

        print(f"Saved: {save_path}")
    

input_dir = "./01_annotation/sample/for_crop"
output_dir = "./01_annotation/sample/cropped"

os.makedirs(output_dir, exist_ok=True)

for file in os.listdir(input_dir):
    if not file.endswith(".txt"):
        continue

    txt_path = os.path.join(input_dir, file)
    img_name = os.path.splitext(file)[0] + ".jpg"
    img_path = os.path.join(input_dir, img_name)

    if not os.path.exists(img_path):
        print(f"Image not found: {img_path}")
        continue

    crop_from_yolo(txt_path, img_path, output_dir)

実行方法

最後に、今回作成したアノテーションツールの実行方法を紹介します。

必要なライブラリ

画像の読み込みにはPillowを使用しています。

Pillowが入っていない場合はインストールします。

pip install pillow

ファイルを配置する

まず、画像とYOLOラベルをfor_cropフォルダへ配置します。

例えば次のような構成です。

01_annotation/
├── sample/
│   ├── for_crop/
│   │   ├── 001.jpg
│   │   ├── 001.txt
│   │   ├── 002.jpg
│   │   ├── 002.txt
│   │   └── ...
│   │
│   └── cropped/
│
└── crop.py

画像とラベルは、拡張子以外を同じ名前にしておきます。

プログラムを実行する

その後、ターミナルから00_crop.pyを実行します。

python crop.py

環境によっては次のように実行します。

python3 crop.py

処理が完了すると、

./01_annotation/sample/cropped

へ切り出した牌画像が保存されます。

まとめ

今回は、前回作成したYOLOラベルを使って、元画像から麻雀牌を自動で切り出しました。

一般的なYOLOの学習では、

撮影
↓
アノテーション
↓
学習

と進めることができます。

しかし今回は、学習データを自動生成するために、あえてその間に**「牌画像を素材として切り出す工程」**を追加しています。

今回の流れでは、

撮影した画像
↓
矩形アノテーション
↓
牌素材を切り出す
↓
柄部分を4点アノテーション
↓
射影変換して別画像へ貼り付ける
↓
学習画像・YOLOラベルを自動生成する

という独自の手順でデータセットを作っていきます。

今回の処理だけを見るとシンプルなクロップ処理ですが、この後の学習データ大量生成につなげるための重要な前処理になっています。

次回は、この切り出した牌画像に対して、射影変換に使用する4点をアノテーションするツールを作ります。

関連記事

コメント

タイトルとURLをコピーしました