Magicode logo
Magicode
7 min read

【AIの画像認識】自作データを簡単にtrainとvalidに分けよう ~学習準備編~

https://cdn.apollon.ai/media/notebox/00234442-f9ca-406d-948b-2b7a4b252c16.jpeg
前回の記事で、AIへの学習方法を記事にしました。
しかし、学習データと検証データを分けることに、めんどくささ感じたのでプログラムで解決しました。
前回の復習

学習に使うフォルダーの中身について

学習のために以下のようなフォルダー構成にする必要がありました。
yolov5
 ├─ data.yaml
 ├─ train
 │    ├─ images
 │    │     └─ *.jpg
 │    ├─ labels
 │          └─ *.txt
 ├─ valid 
       ├─ images
       │     └─ *.jpg
       ├─ labels
             └─ *.txt
data.yamlの中身

train : train/images # 学習画像パス
val: valid/images # 検証用画像パス
nc: 2 # クラス数
names: ["dummy_1" , "dummy_2"] # クラス名

だが、このフォルダー構成には問題があって、多分多くの人は以下のようになると思っています。
train_data
    ├─ classes.txt
    ├─ *.jpg
    ├─ *.txt
labelimgツールはlabelの保存先を自由に選べるのだが、 あまり分け過ぎると混乱の元
だからこそ、同じフォルダーに入れておきたい。
なんなら保存先もあまりいじりたくない。
そして、trainとvalidに分けるのもめんどくさい、data.yamlを書く気も起きない。
そんな問題を一括で処理するようにしました。

アノテーション済み自作データをすぐに学習へ使えるようにするプログラム

python
import glob
import os

import numpy as np
from natsort import natsorted


def main(input_path, data_mode = "w"):
    
    input = input_path.copy()
    output_list = []
    
    train_path = "./train.txt"
    valid_path = "./valid.txt"
    test_path = "./test.txt"
    
    # アノテーション済みのtxtファイルから、画像ファイルを探す
    for input_txt in natsorted(input_path):
        if input_txt[-4:] == ".txt":

            for input_images in natsorted(input):
                if input_images[-4:] == ".txt":

                    flg = False

                # Jpegファイル時にエラー原因の可能性あり
                elif input_txt[:-4] == input_images[:-4] or input_txt[:-4] == input_images[:-5]:

                    flg = True
                    break

                else:
                    flg = False

            if flg:
                input.remove(input_txt)
                input.remove(input_images)

                output_list.append(f"{input_images}\n{input_txt}\n")


    # 中身シャッフル
    rng = np.random.default_rng()
    shuffle = rng.permutation(output_list)

    train_list = shuffle[: int(len(shuffle) * 0.6)]
    valid_list = shuffle[int(len(shuffle) * 0.6) : int(len(shuffle) * 0.8)]
    test_list = shuffle[int(len(shuffle) * 0.8) :]

    with open(train_path, data_mode) as train_txt:

        for output in natsorted(train_list):
            train_txt.write(output)

    with open(valid_path, data_mode) as valid_txt:

        for output in natsorted(valid_list):
            valid_txt.write(output)

    with open(test_path, data_mode) as test_txt:

        for output in natsorted(test_list):
            test_txt.write(output)

    data_yaml = []
    
    # classesからdata.yamlのデータ作成
    for classes in input:

        if "classes" in classes:
            with open(classes, "r") as yaml:
                for data in yaml:
                    data_yaml.append(data[:-1])

    # なぜかここでやらないとうまくいかなかった(コードの書き方に問題があるのかも)
    from ruamel import yaml

    yaml_content = yaml.load(
        """
    train: training
    val: validation
    test: test

    nc: 0
    names: class_name
    """,
        Loader=yaml.Loader,
    )

    yaml_content["train"] = f"{os.path.abspath(train_path)}"
    yaml_content["val"] = f"{os.path.abspath(valid_path)}"
    yaml_content["test"] = f"{os.path.abspath(test_path)}"
    yaml_content["nc"] = len(data_yaml)
    yaml_content["names"] = data_yaml

    # new_yaml = yaml.dump(yaml_content, Dumper=yaml.RoundTripDumper)
    with open("data.yaml", "w") as stream:
        yaml.dump(yaml_content, Dumper=yaml.RoundTripDumper, stream=stream)


if __name__ == "__main__":
  
	input_path = glob.glob(r"C:\Magicode\*")
    main(input_path)
    
    # データセットごとにディレクトリが分けられている場合 
    # pathに上層フォルダの絶対パスを入れてアンコメントしてください  
        
    # path = glob.glob(r"C:\Magicode\*")
    # input_path = glob.glob(r"C:\Magicode\*")
    # for path in input_path:
    #     paths = glob.glob(path + "\*")
    #     main(paths,"a")
使い方は単純
input_pathに学習させたい画像とラベルが入ったファイルの絶対パス(shift + 左clickのパスのコピー)を入力して実行
magicode(上記プログラムの入ったファイル)
    ├─ data.yaml
    ├─ test.txt
    ├─ train.txt
    └─ valid.txt
4つのファイルが作成できていれば完了です。

作成したファイルの中身

今回作成したdata.yamlの中身

train : 絶対パス/train.txt # 学習画像パスへのテキストファイル
val: 絶対パス/valid.txt # 検証用画像パスへのテキストファイル
test: 絶対パス/test.txt # テスト画像パスへのテキストファイル
nc: 2 # classes.txtの中身の数
names:
  • "dummy_1" # classes.txtの中身
  • "dummy_2" # data.yamlのリスト表記

train, valid, testの.txtファイルの中身

train.txt (6割)
絶対パス/*jpg
絶対パス/*txt
valid.txt (2割)
絶対パス/*jpg
絶対パス/*txt
test.txt (2割)
絶対パス/*jpg
絶対パス/*txt

学習開始

C:\>git clone https://github.com/ultralytics/yolov5
C:\>cd yolov5
C:\yolov5>python train.py ……
--data data.yaml --cfg models/hub/yolov5n6.yaml --weights "" --batch-size 8 --epochs 300 data.yamlの部分に生成したdata.yamlを入れる(ないしは絶対パスを入れる)
これで学習が手軽に始めれます。

プログラム及び学習に関する問題点

学習に使用するライブラリprotobufのバージョンが最新バージョンだとエラーが出ます(2022/6/4現在)
C:\>pip uninstall protobuf
C:\>pip install protobuf==3.20
バージョンエラーは上記のコマンドで解決できます。
そして、一番の問題は、プログラムがtxtの名前から画像ファイルの名前を探す方式を取る都合上
JPEGファイルに対応するために以下のような書き方をしているのですが、
elif input_txt[:-4] == input_images[:-4] or input_txt[:-4] == input_images[:-5]:
フォルダーの名前と構成次第ではエラーが引き起こされます。 拡張子に合わせて変更するか、探してくれることにかまけてデタラメな構成をしない限りは、
natsortedによってうまくいくような処理に仕上がっているとは思っています。

最後に、

一番めんどくさいのはこれらのデータわけではなくアノテーションであるという事実から目を逸らすこととします!
という冗談は置いといて、アノテーションに便利なツールは作成中です。
見つけたり、できたりしましたら記事にしますので、その時もぜひよろしくお願いします。

Discussion

コメントにはログインが必要です。