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

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")