[PyTorch] 4-1. 나만의 이미지 데이터셋 만들기

이번 포스트에서는 이전 시간에 다루었던 DataLoader를 가지고 나만의 이미지 데이터 셋을 만들어보는 실습을 진행해보겠습니다.

 

1. 이미지 데이터 파일 만들기

이미지 데이터 셋을 만들기 위해서는 자신이 원하는 이미지 파일이 있어야합니다. jpg, jpeg, png, bmp 등등의 이미지 형식의 파일들을 class별로 폴더를 만들어 넣어줍니다. 예시로써 저는 가방, 신발, 티셔츠 총 3가지 클래스를 만들어서 각각 8장의 이미지 파일을 넣어주었습니다. 여기서 주의해야 될 점이 저의 경우에는 "clothes"라는 상위 폴더를 만들고 그 안에 앞서말한 3가지의 클래스 폴더를 넣어주었습니다. 그렇기 때문에 밑에서 폴더 경로를 가지고 올 때는 상위 폴더의 경로를 가지고와야 합니다.

 

2. 필요 라이브러리 가져오기

필요한 라이브러리들을 import 해줍니다. 이미지 파일을 시각화하기 위해 matplotlib을 사용하겠습니다. 또한 torchvision에서 이미지를 전처리 할 수 있는 transforms을 사용하고 torch.utils.data에서 Dataset과 DataLoader를 사용하여 데이터를 가져와 보겠습니다.

 

3. 데이터 전처리

우선 이미지를 전처리합니다. ransforms.Compose를 통해 Resize, ToTensor, Normalize를 시켜줍니다. 이미지의 높이와 너비가 조금씩 다를 수 있고, 이미지 데이터를 정규화 해주기위해 전처리를 진행합니다. ToTensor를 통해 이미지를 텐서형태로 바꿔줍니다.

그 다음 trainset을 ImageFolder를 사용하여 저장합니다. 경로는 클래스 폴더들이 들어있는 상위폴더로 경로를 지정해줍니다.

또한 앞서서 선언했던 trans를 transform 파라미터 값으로 넣어 데이터가 전처리 될 수 있도록 합니다.

 

trainset에서 __getitem__ 메소드를 통해 18번째 데이터를 가져와 봤습니다. 일부분만 나오지만 0부터 1사이의 값으로 정규화 된 것을 볼 수 있고 맨 끝에는 라벨값인 2가 보입니다.

 

trainset의 총 데이터 갯수를 len()함수를 통해 확인해보았습니다. 3개의 클래스들이 각각 8개씩 사진을 가지고 있기 때문에 총 데이터 수는 24입니다.

trainset에서 classes 메서드를 사용해서 classes 변수에 저장해줍니다. 출력해보면 제가 지정했던 클래스 폴더 명이 그대로 나오는 것을 볼 수 있습니다.

4. DataLoader를 통해 데이터 올리기

이번에는 ImageFolder를 통해 불러온 trainset을 DataLoader를 사용하여 Batch형식으로 네트워크에 올릴 수 있도록 준비해보겠습니다.

DataLoader를 통해 trainset을 감싸주고 batch_size는 16개로 설정했습니다. shuffle은 원래 train 데이터의 경우 섞는게 맞지만 이번 예제에서는 섞지 않도록 하겠습니다. 사용할 CPU core 갯수는 4개로 지정했습니다.

 

iteration을 통해 trainloader에 있는 데이터들을 한 batch 씩 꺼내어 확인해 볼 수 있습니다. labels를 출력해보면 0이 8개 1이 8개 인것을 볼 수 있습니다.

5. 이미지 데이터 시각화

다음으로 불러온 데이터를 시각화 할 수 있는 함수를 만들어보겠습니다. 이전 글에서 설명했던 부분과 동일한 내용입니다.

 

마지막으로 저장한 데이터를 시각화해본 부분입니다. 4번째 출력 부분을 보면 16, 3, 100,100이라고 되어있는데 batch_size가 16이었고 색깔이미지이기 때문에 3입니다. 또한 이미지 사이즈를 100 x 100으로 Resize 해주었기 때문에 16,3,100,100으로 출력됩니다.

 

이상으로 자기만의 이미지 데이터셋을 만들어보았습니다. 이것을 통해서 분류나 이미지 디텍팅 같은 문제들을 자신만의 데이터셋을 활용하여 다양한 시도를 해볼 수 있겠습니다.

 

소스코드 :

https://github.com/Hinterhalter/Deep-Learning-with-PyTorch/blob/master/MyData.ipynb

Hinterhalter/Deep-Learning-with-PyTorch