전처리) train_list.와 val_list 파일 나누기

그동안 Yolo와 같은 모델을 통해 학습하기 위해 정말 많은 전처리를 하게 되었다. 그 중엔 xml 파일을 yolo 스타일의 txt 파일로 변경하거나 json 파일을 똑같이 yolo txt 스타일로 변경하는 일을 자주 하게 되었다. 그동안 한번 코드를 짜둔 후 계속 복사해서 사용하다 보니 꽤나 불편하고 오랜만에 보면 기억도 잘 안나서 블로그에 정리해 두고 깃허브에 코드를 올려둘 계획이다.

아직 yolov5 이하의 모델을 사용해본적이 없어서 다른 모델은 모르겠지만 학습 중에 학습을 하고 그 것을 검증할 데이터 셋을 나눠야 한다.

아래 코드는 그 코드이다.

 

from glob import glob
path = "학습을 돌릴 path" #학습을 돌릴 path => train.txt와 val.txt를 저장할 path
image_path= "your path" # 이미지가 있는 디렉토리, 나는 절대 경로를 넣어두는 편이다.
img_list = glob(f'{image_path}/*') # 이미지 파일들의 이름 들을 읽어온 후 리스트로 저장한다.
 
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)
#test_size => 전체 데이터 셋에서 val_img_list 가 가져갈 비율
# random_state => 섞는 비율
 
print(len(train_img_list),len(val_img_list))
 
with open(f'{path}/train.txt','w') as f:
	f.write('\n'.join(train_img_list)+'\n')
 
with open(f'{path}/val.txt','w') as f:
	f.write('\n'.join(val_img_list)+'\n')