Yolo 8 用のデータセットを自動で作成するpythonコード

Tobilog-Selection

・Laravel 8にてenvファイルが反映されない場合の対処法
・現代錯誤　吾輩はトビウオである。名前はトビログ。
・久々の鉄道模型　SBB 簡易寝台車

こんにちは！とびろぐ管理人のとびうおです！

今回は、yolo8にて、データセットを一つのフォルダに作ったけど、その後に適切なディレクトリに保存したい！という時に使えるスクリプトを用意したので紹介します！

スクリプト
解説
作られるディレクトリ
関数の引数
おわりに

スクリプト

以下のスクリプトを実行することで、自動で振り分けることが可能です！

import os
import shutil
from sklearn.model_selection import train_test_split

def prepare_yolo_dataset(source_dir, dest_dir, test_size=0.2):
    # 画像ファイルとラベルファイルのリストを作成
    images = [f for f in os.listdir(source_dir) if f.endswith('.png')]
    labels = [f.replace('.png', '.txt') for f in images]

    # トレーニングと検証用に分割
    images_train, images_val, labels_train, labels_val = train_test_split(
        images, labels, test_size=test_size, random_state=42)

    # ディレクトリを作成（存在しない場合）
    os.makedirs(os.path.join(dest_dir, 'images', 'train'), exist_ok=True)
    os.makedirs(os.path.join(dest_dir, 'images', 'val'), exist_ok=True)
    os.makedirs(os.path.join(dest_dir, 'labels', 'train'), exist_ok=True)
    os.makedirs(os.path.join(dest_dir, 'labels', 'val'), exist_ok=True)

    # ファイルをコピー
    for img in images_train:
        shutil.copy(os.path.join(source_dir, img), os.path.join(dest_dir, 'images', 'train'))
        shutil.copy(os.path.join(source_dir, img.replace('.png', '.txt')), os.path.join(dest_dir, 'labels', 'train'))

    for img in images_val:
        shutil.copy(os.path.join(source_dir, img), os.path.join(dest_dir, 'images', 'val'))
        shutil.copy(os.path.join(source_dir, img.replace('.png', '.txt')), os.path.join(dest_dir, 'labels', 'val'))

# 使用例
source_directory = ''  # 元のディレクトリパス
destination_directory = ''  # 保存先のディレクトリパス

prepare_yolo_dataset(source_directory, destination_directory)