Fashion-MNIST csv파일로부터 이미지 저장하기

AI/Computer Vision

Fashion-MNIST csv파일로부터 이미지 저장하기

체봄 2021. 2. 11. 20:58

Fashion-MNIST 모델은 keras, tensorflow, pytorch 등 다양한 환경에서 간편하게 다운받을 수 있다.

하지만 내가 필요한 것은 Fashion-MNIST 모델이 아닌 이미지였다.

옷의 모양을 학습시키기 위한 데이터가 필요했는데, Fashion-MNIST의 이미지들은 배경이 제거되고 옷의 정면만 나오도록 처리된 깔끔한 이미지이며 grayscale이기 때문에 옷의 무늬나 색상과 같은 특징을 학습하기보단 옷의 모양을 학습하기에 적합할 것 같다는 생각이었다.

이곳저곳 검색해봤지만 모델 형태로만 존재하고 이미지의 형태로는 따로 존재하지 않는 것 같아서, csv 파일을 찾아 직접 이미지로 변환하는 과정을 수행했다.

dacon.io/competitions/open/235594/data/

[이미지] Fashion MNIST : 의류 클래스 예측

출처 : DACON - Data Science Competition

dacon.io

우선 csv 파일(엑셀 형태의 파일)은 이 사이트에서 다운받을 수 있다. (회원가입 필요)

다운받은 후 train.csv 파일을 확인해보면 각 줄마다 index, label, 784개의 pixel 값이 순서대로 적혀있다.

첫 번째 행에 적혀 있는 'index', 'label', 'pixel-'는 필요가 없어서 삭제해주었다.

그리고 index 열(A열)도 필요가 없어서 삭제해주었고, label 열(B열)도 따로 떼어내서 새로운 label.csv 파일로 저장해주었다.

그러면 이제 train.csv 파일에는 각 행마다 하나의 이미지의 784개 pixel 값들이 저장되어 있다.

이를 이미지 파일로 변환하여 저장하는 과정에는 numpy, opencv 모듈을 사용한다.

import cv2
import numpy as np

train = np.loadtxt('train.csv', delimiter=',', dtype=np.uint8)
label = np.loadtxt('label.csv', delimiter=',', dtype=np.uint8)

dtype을 지정해주지 않으면 float형으로 저장된다.

train.shape는 (60000, 784)이고, label.shape는 (60000, )이다. 28x28 이미지가 60000장 존재하기 때문이다.

# (60000, 784) -> (60000, 28, 28)
train_reshape = train.reshape(60000, 28, 28)

train.shape를 (60000, 784)에서 (60000, 28, 28)로 바꿔주는 과정이다.

이렇게 reshape를 해줘야 하나하나의 이미지에 접근할 수 있다.

이 데이터셋은 총 10가지의 label로 구분되는데, 나의 경우 의류만 필요하기 때문에 0, 1, 2, 3, 4, 6번 label인 이미지만 가져올 것이다. 아래 코드에서 ok_lab이라는 리스트에 저장한다.

ok_lab = [0, 1, 2, 3, 4, 6]
i=1	# 파일명을 다르게 지정하기 위한 변수

for img, lab in zip(train_reshape, label):
  if lab in ok_lab:
    # 256x256 resize, 보간법 적용
    img_resize = cv2.resize(img, dsize=(256,256), interpolation=cv2.INTER_CUBIC)  

    # 선명하게 만들기
    kernel = np.array([[0, -1, 0],
                      [-1, 5, -1],
                      [0, -1, 0]])
    img_sharp = cv2.filter2D(img_resize, -1, kernel)

    # 이미지 파일 생성
    cv2.imwrite(str(lab) + '_' + str(i)+'.jpg', img_sharp)	# 파일명 ex : 0_1.jpg

  i += 1

for문에서는 train_reshape와 label로부터 하나의 이미지 배열과 그에 해당하는 레이블을 하나씩 가져온다.

레이블이 ok_lab에 지정한 의류 레이블인 경우에만 이미지 파일로 저장하는 과정을 수행한다.

먼저 Fashion-MNIST의 기본 이미지가 28x28 크기이기 때문에, 내가 필요한 크기(256x256)로 resize를 한다.

이 때 이미지의 크기를 확대하면 이미지가 깨지기 마련이다. 따라서 interpolation(보간법)을 적용한다. 확대하는 경우에는 INTER_CUBIC이나 INTER_LINEAR를 자주 사용한다고 하는데, 결과물을 확인해보았을 때 INTER_CUBIC이 더 나은 결과를 보여서 사용해주었다.

보간법을 적용해도 화질이 너무 좋지 않아서 선명하게 만드는 작업을 추가하였다.

이제 imwrite() 함수로 이미지 파일을 생성한다. 인자에는 파일명과 배열을 넣어주면 된다.