3. Image preprocessing 및 train, val 나누기

본 내용에서는 손가락의 숫자를 실시간으로 맞추는 것을 목표로 한다.

 

본 내용은 window 환경에서 진행하며, Anaconda를 깔아둔 상태이다. 만약 Anaconda 환경이 만들어져 있지 않다면 아래 링크를 참조합니다.

https://silinu-ai.tistory.com/3

Jupyter notebook(주피터 노트북) 설치하기

 

이전 글에서는 데이터 수집, 이미지 라벨링을 진행했다. 

이전 글은 아래 링크로 들어갈 수 있다.

• 데이터 수집
• 이미지 라벨링

 

이번에는 이미지와 라벨을 학습할 데이터(train data)와 모델 검증을 위한 데이터(validation data)로 나눈다.

train data와 validation data은 주로 7:3이나 8:2의 비율을 사용한다.

본 내용에서는 8:2의 비율을 사용한다.

 

이제 다음과 같이 진행한다.

1. 이미지 저장 폴더에서 이미지 리스트 불러오기
2. 데이터를 임의로 섞은 후, train data와 validation data를 나누기
3. 이미지 가공 후 저장하기

 

본 코드의 위치는 편의상 A 폴더라고 지칭한다.

import 할 모듈은 다음과 같다.

import os
import cv2
import yaml
import random
import numpy as np

 

만약 cv2가 없다면 anaconda prompt 창의 가상 환경에서 다음과 같이 opencv를 설치한다.

conda install opencv

 

1. 이미지 저장 폴더에서 이미지 리스트 불러오기 

이전 포스트 보면 이미지 저장 폴더는 A 폴더/All_Data 폴더에 위치한다.

All_Data 폴더에 있는 리스트를 전부 가져온다.

파일의 이름을 가져오기 위해서 .png로 자른 후 파일의 이름을 가져온다.

path = os.getcwd()
Img_path = os.path.join(path,"All_Data")
file_list = os.listdir(Img_path)
file_name_list = [i.split(".png")[0] for i in file_list]

 

2. 데이터를 임의로 섞은 후, train data와 validation data를 나누기

file_name_list 변수에는 데이터 이름 순서대로 저장되어 있다.

이를 shuffle 한 후 8:2의 비율로 나눠 train 값과 validation으로 나눈다.

random.shuffle(file_name_list)
Split_Num = int(len(file_name_list) * 0.8)
Train_file_name = file_name_list[:Split_Num]
Valid_file_name = file_name_list[Split_Num:]

 

3. 이미지 가공 후 저장하기

먼저 저장할 위치를 만들어준다.

• 현재 저장된 이미지 저장 위치: A 폴더/All_Data/
• 현재 저장된 라벨 저장 위치: A 폴더/All_Data_label/
• train data의 image data 저장할 위치: A 폴더/train/images/
• train data의 label data 저장할 위치: A 폴더/train/labels/
• validation data의 image data 저장할 위치: A 폴더/valid/images/
• validation data의 label data 저장할 위치: A 폴더/valid/labels/

path = os.getcwd()

pre_img_file_path = os.path.join(path,"All_Data")
pre_label_file_path = os.path.join(path,"All_Data_label")

Train_path = os.path.join(path,"train")
Valid_path = os.path.join(path,"valid")

Train_img_path = os.path.join(Train_path,"images")
Train_label_path = os.path.join(Train_path,"labels")
Valid_img_path = os.path.join(Valid_path,"images")
Valid_label_path = os.path.join(Valid_path,"labels")

path_list = [Train_path,Valid_path,Train_img_path,Train_label_path,Valid_img_path,Valid_label_path]

for p in path_list:
    if not os.path.isdir(p):
        os.mkdir(p)

 

한글이 포함된 경로를 읽고 저장한다면 다음 함수를 작성한다.

# 파일 위치가 한글 경로가 존재할 때 사용함(없다면 제외하고 해도 상관 없음)
# 출처 : https://jangjy.tistory.com/337
def imread(filename, flags=cv2.IMREAD_COLOR, dtype=np.uint8):
    try:
        n = np.fromfile(filename, dtype)
        img = cv2.imdecode(n, flags)
        return img
    except Exception as e:
        print(e)
        return None

# 파일 위치가 한글 경로가 존재할 때 사용함(없다면 제외하고 해도 상관 없음)
# 출처 : https://jangjy.tistory.com/337
def imwrite(filename, img, params=None):
    try:
        ext = os.path.splitext(filename)[1]
        result, n = cv2.imencode(ext, img, params)
        if result:
            with open(filename, mode='w+b') as f:
                n.tofile(f)
            return True
        else:
            return False
    except Exception as e:
        print(e)
        return False

 

이미지를 가공하는 함수이다.

한글 경로로 되어 있지 않다면 1번의 주석을 풀고 2번을 주석처리한다.

한글 경로로 되어있다면 그대로 2번 코드를 사용한다.

여기서는 사이즈 조절만 하였지만 모자이크, 흑백 등의 전처리를 추가하고 싶다면 이 함수에서 작성하면 된다.

def img_preprocessing(img_path):
    # 한글 경로로 되어 있지 않을 때 1번, 한글 경로로 되어 있다면 2번을 사용함.
    # 1번
    #re_img = cv2.imread(img_path)
    # 2번
    img = imread(img_path)
    
    #전처리 부분
    re_img = cv2.resize(img, (416, 416))
    return re_img

 

이미지와 라벨을 저장하는 함수이다.

위와 마찬가지로 저장할 이미지 경로에 한글이 없다면 1번 주석 풀고 2번을 주석처리한다.

반대로 한글이 경로에 포함된다면 그대로 2번 코드를 사용한다.

def save_files(file_name, Img_path, Change_img_path, Label_path, Change_label_path):
    for file in file_name:
        # 이미지 읽고 전처리 후 저장
        img_path = os.path.join(Img_path,file+".png")
        T_img_path = os.path.join(Change_img_path,file+".png")
        p_img = img_preprocessing(img_path)
        
        # 저장할 이미지 경로에 한글이 없다면 1번, 있다면 2번을 실행한다.
        # 1번
        #cv2.imwrite(Train_img_path,p_img)
        # 2번
        imwrite(T_img_path,p_img)

        # 라벨 저장하기
        label_path = os.path.join(Label_path, file+".txt")
        T_label_path = os.path.join(Change_label_path, file+".txt")
        f = open(label_path, 'r')
        h = open(T_label_path,'w')
        data = f.read()
        h.write(data)
        f.close()
        h.close()

 

파일 저장 함수를 실행한다.

# 파일 저장
save_files(Train_file_name, pre_img_file_path, Train_img_path, pre_label_file_path, Train_label_path)    
save_files(Valid_file_name, pre_img_file_path, Valid_img_path, pre_label_file_path, Valid_label_path)

 

class 정보가 담긴 yaml 데이터(data.yaml)를 A 폴더에 저장한다.

# train 위치, valid 위치 설정하기
# nc = classes num(여기서는 0~5까지 6개이다.)
# names = class 이름

with open(os.path.join(pre_label_file_path,"classes.txt"),'r') as f:
    data = f.read()
    classes = data
names = classes.split("\n")[:-1]
nc = len(names)
yaml_data = {'train': Train_img_path, 'val': Valid_img_path, 
          'nc': nc, 'names': names}
yaml_path = os.path.join(path,'data.yaml')
with open(yaml_path,'w') as f:
    yaml.dump(yaml_data,f)
# 잘 저장되었는 지 확인하기
with open(yaml_path,'r') as f:
    yaml_file = yaml.load(f, Loader=yaml.FullLoader)
    display(yaml_file)