[객체인식]train_val 나누기

#train test set 분리

import os

path = "dataset/images/train/"
img_list_path = os.listdir(path)
img_list = [file for file in img_list_path if file.endswith(".jpg")]
print(img_list)
print("\n\n")

# from glob import glob # 파일 경로명까지 저장 => 폐기

# img_list = glob('dataset/images/train/*.jpg')
# print(len(img_list))
# print(img_list)


#train test set 분리
from sklearn.model_selection import train_test_split

train_img_list, val_img_list = train_test_split(img_list, test_size=0.2, random_state=2000)

print(len(train_img_list), len(val_img_list))
print("\n\n")

# val에 넣고 train에 넣을까

import shutil 

# img파일 train val으로 나누기
for a in val_img_list: 
    filename = a
    src = 'dataset/images/train/' 
    dir = 'dataset/images/val/' 
    shutil.move(src + filename, dir + filename)
print(len(val_img_list),val_img_list)
print("\n\n")

# 나뉜 img_val을 기준으로 labeling된 txt 파일  train val 으로 나누기
val_txt_list=[]
for a in val_img_list:
    temp = a.replace("jpg","txt")
    val_txt_list.append(temp)
print(len(val_txt_list),val_txt_list)
print("\n\n")

for a in val_txt_list: 
    filename = a
    src = 'dataset/labels/train/' 
    dir = 'dataset/labels/val/' 
    shutil.move(src + filename, dir + filename)