プログラミングの芽

面白そうな技術を使って、楽しいことやりたい

機械学習に使う画像とアノテーションファイルをtestとtrainフォルダに振り分ける

2500枚とそれに付随するアノテーションファイルを良い感じに振り分けることを目的とします。

必要なもの

import pandas as pd
import numpy as np
import cv2
from sklearn.model_selection import train_test_split


アノテーションファイルをtrain、test用にランダムに振り分ける

df = pd.read_csv('coin_labels.csv')
namelist_train, namelist_test = train_test_split(df, stratify=df['class'])

train_test_split()の第二引数は今回だとclassにしてます。
こうするとクラスの分布を良い感じに散らばして分けてくれるみたいです。
参考
scikit-learn でトレーニングデータとテストデータを作成する – Python でデータサイエンス

0列目を消してそれぞれをcsvに保存

namelist_train.to_csv("train.csv",index=False)
namelist_test.to_csv("test.csv",index=False)


画像の振り分けをします。
変数やら準備

dir_path = "./coin/"

train_img_name = np.array(namelist_train["filename"],dtype="unicode")
test_img_name = np.array(namelist_test["filename"],dtype="unicode")

train_cnt = train_img_name.shape[0]
test_cnt = test_img_name.shape[0]

train_img_path = []
test_img_path = []


train、testフォルダに振り分ける

for i in range(train_cnt):
    train_img_path.append(dir_path+train_img_name[i])
    train_img = cv2.imread(train_img_path[i])
    cv2.imwrite("./train/{0}".format(train_img_name[i]), train_img)
    
for i in range(test_cnt):
    test_img_path.append(dir_path+test_img_name[i])
    test_img = cv2.imread(test_img_path[i])
    cv2.imwrite("./test/{0}".format(test_img_name[i]), test_img)


以上です。
プログラミングのセンスがない人って、変数をむやみやたらに使いたがる上に変数名もよくわからないみたいですね。私ですね。