【PythonでYOLO自動化】面倒なデータセット準備を秒殺！3つの便利スクリプトを公開

こんにちは、Tech Samuraiです！
AIによる物体検出モデル「YOLO」を使ったカスタム学習は、非常にエキサイティングですよね。しかし、YOLOの学習を経験した人なら誰もが同意する、一つの「壁」があります。それは、モデルのチューニングではなく、その前段階にある**膨大で、地味で、面倒な「データセット準備」**の作業です。

アノテーションした大量の画像とラベルファイルを、手作業で「train（訓練用）」と「val（検証用）」に振り分け、クラス番号を書き換え…。考えただけで気が遠くなります。

今回の記事では、このYOLO学習で最も面倒な前処理作業を自動化するために私が作成した、**3つのPythonスクリプト**を紹介します。これらのツールを使えば、これまで何時間もかかっていた作業が、文字通り数秒で完了します！

前提：アノテーション作業は完了していますか？

まず、非常に重要な前提です。この記事で紹介するスクリプトは、アノテーション作業そのものではなく、**アノテーション済みの画像とラベルファイル**を「YOLOが学習できる形式」に整理・整形するためのものです。

まだアノテーション作業が完了していない方は、まずはlabelImgのような優れたツールを使って、学習用のデータ（画像と、それに対応するYOLO形式の.txtラベルファイル）を作成する必要があります。

labelImgの使い方がわからない方は、例えば、以下のnoteの記事などが非常に分かりやすく解説されています。
参考： 【YOLO v5】labelImgのインストールと使い方（Windows）

ステップ1：データセットの「お掃除」スクリプト

スクリプト: 0_clear_dataset.py

目的： 新しい学習を始める前に、古いデータをリセットする。

YOLOの学習では、パラメータを変えて何度も実験を繰り返します。その際、dataset/images/train`などのフォルダに古いデータが残っていると、新しいデータと混ざってしまい、学習結果が汚染されてしまいます。

このスクリプトは、YOLOの標準的なディレクトリ構造（images/train, images/val, labels/train, labels/val`）を保ったまま、中身のファイルだけを安全に一括削除してくれる「リセットボタン」です。これで、いつでもクリーンな状態で新しい実験を始められます。

ステップ2：データの「自動振り分け」スクリプト

スクリプト: 1_preparation_dataset.py

目的： アノテーション済みデータを、訓練用と検証用に自動でランダムに振り分ける。

これこそが、このプロジェクトの心臓部です。手作業でやると最も時間がかかる、「画像と、それに対応するラベルファイルを、正しいペアのまま、8:2の割合でランダムに`train`と`val`フォルダに振り分ける」という面倒な作業を、一瞬で実行します。

# スクリプトの心臓部（1_preparation_dataset.py より抜粋）

# --- 全ての画像ファイルのリストアップ ---
image_paths = []
for ext in ('*.jpg', '*.png', 'JPG'):
    image_paths.extend(glob.glob(os.path.join(source_dir, ext)))

# --- 画像ファイルリストをシャッフル ---
random.shuffle(image_paths)

# --- 訓練用と検証用に分割 (8:2の比率) ---
split_index = int(len(image_paths) * train_ratio)
train_image_paths = image_paths[:split_index]
valid_image_paths = image_paths[split_index:]

# --- ファイルを移動する ---
# (この後、各画像に対応する .txt ラベルファイルも一緒に移動する処理が続く)
move_files(train_image_paths, train_img_dir, train_label_dir)
move_files(valid_image_paths, valid_img_dir, valid_label_dir)

ステップ3：クラス番号の「一括置換」スクリプト

スクリプト: 2_modify_labels_recursive.py

目的： ラベルファイル内のクラス番号（インデックス）を一括で修正する。

これは、YOLO学習でよくある「罠」を解決するスクリプトです。例えば、labelImgで「人(0)」「車(1)」「犬(15)」…のように16クラス用のアノテーションをしていたとします。その後、「犬だけを学習させたい」と思っても、ラベルファイルには15`と記録されています。

YOLOで単一のクラスを学習させる場合、そのクラス番号は0`でなければなりません。このスクリプトは、dataset/labels`フォルダ内の全ての.txt`ファイルを再帰的に検索し、指定した番号（例: `15`）を、新しい番号（例: `0`）に一括で置換してくれます。

# スクリプトの心臓部（2_modify_labels_recursive.py より抜粋）

# ...
for line in lines:
    parts = line.strip().split()
    if parts and parts[0] == '15': # ← 変更したい元のクラス番号
        parts[0] = '0' # ← 変更後のクラス番号
        modified_lines.append(' '.join(parts) + '\n')
    else:
        modified_lines.append(line)
# ...