Train/Validation 나눠주기

yolov5 학습을 위해 train : validation을 8:2로 나누어주기 위한 코드. 경로를 작성해줄 때 glob 라이브러리를 이용하던데 라이브러리를 사용하여 실행하면 마지막 파일 앞 경로 부분이 '/'가 아닌 '\\' 이런 식으로 작성이 되어 경로 오류가 생겼다. 따라서 os.listdir과 os.path.join을 이용하여 경로를 합쳐주었다.

import os
from pathlib import Path
import shutil
from sklearn.model_selection import train_test_split

#getting list of images
image_files = os.listdir("이미지경로")
img_path = []
for file in image_files:
    image = os.path.join("이미지경로", file)
    img_path.append(image)
    
img_path = [name.replace(".jpg", "") for name in img_path]

train_names, val_names = train_test_split(
    img_path, test_size=0.2, train_size=0.8, random_state=42, shuffle=True)

# separate train from validation
def batch_move_files(file_list, source_path, label_path, destination_path_l, destination_path):
    for file in file_list:
        image = file  + '.jpg'
        txt = file.split("/")[-1] + '.txt' #yolo라 txt 파일로 함. 
        shutil.copy(os.path.join(source_path, image), destination_path) #images -> train, valid
        shutil.copy(os.path.join(label_path, txt), destination_path_l) # label -> train_l, valid_l
    # return

#image path, label path
source_dir = "이미지 경로"
label_dir = "라벨 경로"

#train image path, train image label path
train_dir = "train image 경로"
trainLa_dir = "train label 경로"

#validation image path, validation image label path
val_dir = "valid image 경로"
valLa_dir = "trian label 경로"

batch_move_files(train_names, source_dir, label_dir, trainLa_dir, train_dir)
batch_move_files(val_names, source_dir, label_dir, valLa_dir, val_dir)

 

 

참고 사이트:

https://lynnshin.tistory.com/46

https://eehoeskrap.tistory.com/496