본문 바로가기

Artificial Intelligence/Tensorflow

[Tensorflow] shutil을 사용한 데이터 분류

728x90

shutil을 사용하여 kaggle에서 제공하는 Food Images (Food-101) 데이터를 학습과 테스트로 나누는 코드입니다. 

 

import os
import shutil

directory_list = ['./food101/images/test_set/', './food101/images/train_set/']

for path in directory_list:
    os.mkdir(path)

 

학습 데이터와 테스트 데이터를 넣어줄 test_set, train_set 폴더를 food101/images 밑에 생성해주는 코드입니다.

 

file = open('./food101/meta/meta/classes.txt')
classes = file.readline()

while(classes):
    path = directory_list[0] + classes[:-1]
    os.mkdir(path)
    classes = file.readline()

 

클래스 이름이 저장되어 있는 classes.txt 한 줄씩 읽어와서 test_set 밑에 classes [:-1] 이름을 가지고 있는 폴더들을 만들어줍니다.

 

test_file = open('./food101/meta/meta/test.txt')
img_path = test_file.readline().strip()

while(img_path):
    classes = img_path.split('/')[0]
    from_path = './food101/images/{}'.format(img_path + '.jpg')
    to_path = directory_list[0] + classes
    shutil.move(from_path, to_path)
    img_path = test_file.readline().strip()

 

테스트 데이터 경로가 담겨져있는 test.txt를 읽어와서 기존에 저장되어 있는 폴더에서 test_set 폴더 밑에 classes 이름에 맞는 클래스 폴더로 옮겨줍니다.

 

file = open('./food101/meta/meta/classes.txt')
classes = file.readline()

while(classes):
    path = directory_list[1] + classes[:-1]
    os.mkdir(path)
    classes = file.readline()

 

테스트 데이터와 마찬가지로 클래스 이름이 저장되어 있는 classes.txt 한 줄씩 읽어와서 train_set 밑에 classes [:-1] 이름을 가지고 있는 폴더들을 만들어줍니다.

 

train_file = open('./food101/meta/meta/train.txt')
img_path = train_file.readline().strip()

while(img_path):
    classes = img_path.split('/')[0]
    from_path = './food101/images/{}'.format(img_path + '.jpg')
    to_path = directory_list[1] + classes
    shutil.move(from_path, to_path)
    img_path = train_file.readline().strip()

 

학습 데이터 경로가 담겨져있는 train.txt를 읽어와서 기존에 저장되어 있는 폴더에서 train_set 폴더 밑에 classes 이름에 맞는 클래스 폴더로 옮겨줍니다.

 

 

file = open('./food101/meta/meta/classes.txt')
classes = file.readline()

while(classes):
    path = './food101/images/{}'.format(classes[:-1])
    shutil.rmtree(path)
    classes = file.readline()

 

학습 데이터와 테스트 데이터를 모두 분류하였으므로 기존 데이터에 들어있던 클래스 이름을 폴더들을 전부 삭제해줍니다.

728x90