[pandas] csv 파일에 있는 정보로 이미지 데이터 load하기

데이콘이나 캐글에서 이미지 데이터를 사용할 때 데이터셋 폴더 안에 csv 파일과 이미지 파일을 같이 주는데 csv 파일 경우에 이미지 파일명이 있는 상대 경로와 label이 명시되어 있다.
그래서 이미지 데이터를 사용하기 위해 csv 파일에 있는 이미지 파일 경로로 이미지를 load하고 label를 달아줘야 한다.

Import

import numpy as np
import pandas as pd
from PIL import Image

PIL을 설치하기 위해서는 PIL이 아닌 pip install Pillow로 설치해줘야한다.

Code

csv 파일 로드

train_path = './data/train/'
test_path = './data/test/'

train_df = pd.read_csv(train_path + 'train_data.csv')
test_csv = pd.read_csv(test_path + 'test_data.csv')

학습 데이터와 테스트 데이터는 프로젝트 root 디렉토리에서 data 폴더 아래에 전부 담겨있다. 근데 학습 데이터는 train 디렉토리 안에, 테스트 데이터는 test 디렉토리 안에 따로 있어 경로명을 분리해줬다.

csv 파일 예시는 위 그림과 같다.

데이터 로드 및 변환

# 리스트 초기화
train_images = []
train_labels = []
test_images = []

# csv에서 데이터 로드하기
for file in train_df['filen_name']:
    image = np.array(Image.open(train_path + file))
    train_images.append(image)
for label in train_df['label']:
    train_labels.append(label)

for file in test_df['file_name']:
    image = np.array(Image.open(test_path + file))
    test_images.append(image)

# list -> ndarray로 변환
train_images = np.array(train_images)
train_labels = np.array(train_labels)
test_images = np.array(test_images)

# 정규화 후 ndarray 형변환
train_images = train_images / 255.
train_images = train_images.reshape(-1, train_images.shape[1], train_images.shape[2], 1)
test_images = test_images / 255.
test_images = test_images.reshape(-1, test_images.shape[1], test_images.shape[2], 1)

이미지 파일도 path 아래에 있기 때문에 반복문을 통해 파일명에 일치하는 이미지를 PIL로 load해 변수 image에 저장하고 변수 image를 리스트 images에 append()한다.
label로 마찬가지로 열 [‘label’]에 있는 행 값을 순서대로 리스트 labels에 넣어주면 이미지 파일과 label이 순서가 맞게 저장된다.
그런 다음에 list를 ndarray로 변환해주고 0-1 사이 값으로 정규화한 다음에 흑백 이미지이기 때문에 reshape()를 수행할 때 채널 수를 1로 해야 하기 때문에 (데이터 개수, width, height, 1)로 변환시킨다.

데이터 저장 및 불러오기 (npy 파일)

그런데 이미지 파일이 많다면 위 방법으로 load하는 것은 많은 시간이 소요되기 때문에 비효율적이다.
그래서 가볍게 하고자 npy 파일로 만들어 load하자.

data_path = 'data/npy/'

# save
np.save(data_path + 'train_images.npy', train_images)

저장할 path를 지정하고 np.save()를 통해 이미지 데이터를 npy 파일로 저장하면 된다.
반대로 저장된 npy 파일을 load하는 방법은 다음과 같다.

train_images = np.load(data_path + 'train_images.npy')