【Deepl Learning】学習リストの作成

Keras-yolo3での学習に使用される学習リストの作成備忘録です。
まだ、学習のスクリプトが決まったのが出来てないので、部品的なスクリプトの小出しです。

以下、スクリプトです。基本的にパスを変えて動作させます。
例としてフルーツの画像判定を使ってます。種類と品種によって異なったモデルの学習リストを作成します。

#!/usr/bin/env python
# coding: utf-8

# # フォルダパスから指定した長さの学習リストを作成する

# ## モジュールインポート

import os, glob, random


# ## 定数宣言

target_path = '../train_set/追加データ(keras-yolo3)' #学習対象フォルダ
train_name = 'train.txt' #学習リスト名

target_type = '品種'
classes_path = 'model_data/品種モデル/6_桃' #クラスファイル名

output_path = '学習リスト'

list_number = 100


# ## 学習リスト作成

if target_type == "品種":
    output_dir = os.path.join(output_path, os.path.basename(target_path), target_type, os.path.basename(classes_path), str(list_number))
else:
    output_dir = os.path.join(output_path, os.path.basename(target_path), target_type, str(list_number))

if not os.path.isdir(output_dir):
    os.makedirs(output_dir)
    print("folder made:", output_dir)
    
class_names = []
classes_file = os.path.join(classes_path, "class.txt")
with open(classes_file, 'r', encoding='UTF-8') as f:
    line =  f.readline()
    while line:
        class_names.append(line.split("\n")[0])
        line =  f.readline()    

target_dir = os.path.join(target_path, target_type, "image")

if target_type == "品種":
    target_dir = os.path.join(target_dir, os.path.basename(classes_path))

if not os.path.isdir(target_dir):
    print("folder does not exist:", target_dir)
    exit(1)

train_list = []
for img_path in glob.glob(os.path.join(target_dir, '*')):
    
    if not os.path.isdir(img_path):
        continue

    class_name = os.path.basename(img_path).split("_")[1]
    class_index = class_names.index(class_name)

    count = 0

    image_list = glob.glob(os.path.join(img_path, '**', '*.jpg'), recursive=True)
    random.shuffle(image_list)
    for img_file in image_list[:list_number]:

        annotation_file = img_file.replace('image', 'annotation').replace('.jpg', '.txt')
        if not os.path.isfile(annotation_file):
            continue

        with open(annotation_file, 'r', encoding='UTF-8') as f:
            # Keras-yolo3のフォーマットに合わせる
            # x_min, y_min, x_max, y_maxの順
            line =  f.readline()
            annotation_list = []
            while line:
                annotation = line.split("\n")[0]
                annotation_list.append(annotation)
                line = f.readline()
            train_list.append([img_file, annotation_list])

        count += 1

        
print("train_list:",len(train_list))

train_file = os.path.join(output_dir, train_name)
with open(train_file, 'w', encoding='utf-8') as f:
    for line in train_list:
        jpg = line[0]
        last = line[1][-1]
        f.write(jpg+' ')
        for annotation in line[1]:
            f.write(annotation)
            f.write(' ')
        f.write('\n')
print("train file made:",train_file)

最近忙しいので、部品スクリプトを小出しする感じになりそうです。