[#04] AI Hub 한국어 글자체 AI 이미지 데이터 전처리

OCR 모델을 이용하여 약국이나 편의점에서 살 수 있는 일반의약품의 상품명을 인식해보는 과정을 담아보는 글.

지난 글에서는 정부에서 관리하는 AI Hub 사이트에서 한글 이미지 데이터셋을 구할 수 있었다. 하지만 태깅 데이터 정보가 담겨있는 json 파일은 262Mb의 크기로 대용량이기 때문에 우리가 원하는 정보만 추출해야하며 또 OCR 모델을 학습시키기 위해 전처리를 할 필요가 있다.

 

데이터를 가공해 학습에 사용하자

데이터 분석

AI Hub에서 제공하는 Text in the Wild 데이터셋의 태깅데이터 textinthewild_data_info.json 파일은 크게 info, images, annotations, licenses 라는 key들로 이루어져있다.

import json
file = json.load(open('./textinthewild_data_info.json'))
file.keys() #dict_keys(['info', 'images', 'annotations', 'licenses'])
file['info'] #{'name': 'Text in the wild Dataset', 'date_created': '2019-10-14 04:31:48'}
type(file['images']) #list

images key에는 모든 이미지 정보가 담겨 있으며 annotations과 공유하고 있는 id와 파일 이름(file_name), 어떤 종류의 이미지인지(book, goods, signboard, traffic sign) 리스트 형태로 저장되어 있다.

file['images'][:3]

'''
[{'id': '00000001',
  'width': 1920,
  'height': 1440,
  'file_name': 'FFF2F34A08347075F55E72B240EFE691.jpg',
  'type': 'book'},
 {'id': '00000002',
  'width': 1920,
  'height': 1440,
  'file_name': 'FFDE6BAEADC9EDD31E51A1D1F687310F.jpg',
  'type': 'book'},
 {'id': '00000003',
  'width': 1920,
  'height': 1440,
  'file_name': 'FFCFEB9E1D09544D6B458717DB4D6B7C.jpg',
  'type': 'book'}]
'''

 

전체 태깅 데이터 중 내가 사용하고자 하는 상품(goods)이미지만 골라보았는데 개수가 26,358개라고 나왔다... 아니, 다운 받은 상품 이미지 폴더 안에 존재하는 파일의 개수는 37,200개였는데...?

근데 이후에 데이터 가공을 더 하다보니까 태깅도 잘못 되어있는 부분들도 많고 해서 뭔가 태깅 누락이 발생한 것 같다는 생각이 들었다. 정부가 관리하는 데이터지만 깔끔하지 않다는거 ^^ 

file['images'][0]['type'] == 'book' # True
goods = [f for f in file['images'] if f['type']=='product']
len(goods) #26358

 

goods[0]

'''
{'id': '00001128',
 'width': 1920,
 'height': 1441,
 'file_name': 'EB862D328745738A1072A0E28F2E444E.jpg',
 'type': 'product'}
'''

첫 번째로 저장된 상품의 태깅 정보를 가져와 보았다. 여기서 주의해야할 점... images 값에 존재하는 id는 annotations 값에서 id가 아닌 image_id와 같은 정보다! 헷갈려 하면 안된다. 그리고 AI Hub 데이터는 친절(?)하게도 이미지에서 글자 단위와 단어 단위로 각각 태깅이 되어있다. 따라서 단어 단위의 태깅 데이터만 필요한 경우에는 annotations 값에서 attributes['class'] == 'word'로 골라낼 수 있다.

annotation = [a for a in file['annotations'] if a['image_id'] == goods[0]['id'] and a['attributes']['class']=='word']
annotation

'''
[{'id': '00029855',
  'image_id': '00001128',
  'text': '페리덱스',
  'attributes': {'class': 'word'},
  'bbox': [455, 610, 1155, 323]},
 {'id': '00029856',
  'image_id': '00001128',
  'text': '혓바늘,',
  'attributes': {'class': 'word'},
  'bbox': [693, 983, 185, 74]},
 {'id': '00029858',
  'image_id': '00001128',
  'text': '입안이',
  'attributes': {'class': 'word'},
  'bbox': [1160, 967, 169, 77]},
 {'id': '00029859',
  'image_id': '00001128',
  'text': '헐었을',
  'attributes': {'class': 'word'},
  'bbox': [1343, 963, 173, 72]},
 {'id': '00029860',
  'image_id': '00001128',
  'text': '때',
  'attributes': {'class': 'word'},
  'bbox': [1537, 958, 53, 72]}]
'''

위 태깅 정보를 갖고 있는 이미지 사진을 출력해보면 아래와 같다.

import matplotlib.pyplot as plt
import cv2
img = cv2.imread('/data/ljh/dataset/ocr/Goods/'+goods[0]['file_name'])
plt.imshow(img)

 

데이터 1차 가공 - AI Hub 데이터 분할

우선 태깅 데이터가 용량이 너무 크고 전체 데이터 정보를 담고 있으므로 상품 이미지에 대해서만 데이터를 추려보도록 하자. 학습에 사용하기 위해 상품 이미지를 train, validation, test set으로 70:15:15의 비율로 나눠서 임의의 순서로 이미지를 나누고 각 이미지에 해당하는 annotation 정보를 함께 저장하였다.

import random
import os

ocr_good_files = os.listdir('/data/ocr/Goods/')
len(ocr_good_files) # 37220

random.shuffle(ocr_good_files)

n_train = int(len(ocr_good_files) * 0.7)
n_validation = int(len(ocr_good_files) * 0.15)
n_test = int(len(ocr_good_files) * 0.15)

print(n_train, n_validation, n_test) # 26054 5583 5583

train_files = ocr_good_files[:n_train]
validation_files = ocr_good_files[n_train: n_train+n_validation]
test_files = ocr_good_files[-n_test:]

## train/validation/test 이미지들에 해당하는 id 값을 저장

train_img_ids = {}
validation_img_ids = {}
test_img_ids = {}

for image in file['images']:
    if image['file_name'] in train_files:
        train_img_ids[image['file_name']] = image['id']
    elif image['file_name'] in validation_files:
        validation_img_ids[image['file_name']] = image['id']
    elif image['file_name'] in test_files:
        test_img_ids[image['file_name']] = image['id']

## train/validation/test 이미지들에 해당하는 annotation 들을 저장

train_annotations = {f:[] for f in train_img_ids.keys()}
validation_annotations = {f:[] for f in validation_img_ids.keys()}
test_annotations = {f:[] for f in test_img_ids.keys()}

train_ids_img = {train_img_ids[id_]:id_ for id_ in train_img_ids}
validation_ids_img = {validation_img_ids[id_]:id_ for id_ in validation_img_ids}
test_ids_img = {test_img_ids[id_]:id_ for id_ in test_img_ids}

for idx, annotation in enumerate(file['annotations']):
    if idx % 5000 == 0:
        print(idx,'/',len(file['annotations']),'processed')
    if annotation['attributes']['class'] != 'word':
        continue
    if annotation['image_id'] in train_ids_img:
        train_annotations[train_ids_img[annotation['image_id']]].append(annotation)
    elif annotation['image_id'] in validation_ids_img:
        validation_annotations[validation_ids_img[annotation['image_id']]].append(annotation)
    elif annotation['image_id'] in test_ids_img:
        test_annotations[test_ids_img[annotation['image_id']]].append(annotation)

with open('train_annotation.json', 'w') as file:
    json.dump(train_annotations, file)
with open('validation_annotation.json', 'w') as file:
    json.dump(validation_annotations, file)
with open('test_annotation.json', 'w') as file:
    json.dump(test_annotations, file)

그러면 아래와 같이 train, validation, test에 쓰이는 상품 이미지들의 파일 이름과 annotation 정보가 담긴 json 파일이 민들어지게 된다.

 

데이터 2차 가공 - Preprocessing

이제 이 태깅 데이터를 맨 처음에 보았던 IC15 데이터셋과 같은 형태로 가공할 필요가 있다. 내가 사용할 Text Recognition 모델은 Naver Clova AI에서 발표한 모델로 github에 학습 데이터 가이드라인이 친절하게 제시되어 있다. 또 블로그 글을 참고하여 데이터 가공을 진행하였다.

네이버 deep-text-recognition 모델을 custom data로 학습 & 아키텍쳐 분석

위와 같은 방식으로 각 이미지에 해당하는 annotation 값을 이용해 'bbox'위치 정보로 단어 영역을 자르고 'text'정보와 함께 저장하였다.

여기서! 태깅된 정보들 중에 잘못된 값들이 상당히 존재했다. bounding box에서 좌상단 좌표를 의미하는 x,y와 너비 및 높이를 의미하는 w,h 값이 0이거나 음수인 경우가 있다는 것을 학습 중에 오류가 발생해 디버깅을 하다가 발견했다. 따라서 처음 데이터를 저장할 때 예외처리를 진행하여 걸러내는 것이 필요하다.

import json
import os
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm

## aihub 데이터 annotation을 읽어서 단어 단위로 잘라서 data에 저장하기

data_root_path = '/data/ljh/dataset/ocr/Goods/'
save_root_path = '/data/ljh/OCR/deep-text-recognition-benchmark/data/'

test_annotations = json.load(open('./test_annotation.json'))
gt_file = open(save_root_path+'gt_test.txt', 'w')
for file_name in tqdm(test_annotations):
    annotations = test_annotations[file_name]
    image = cv2.imread(data_root_path+file_name)
    for idx, annotation in enumerate(annotations):
        x,y,w,h = annotation['bbox']
        if x<= 0 or y<= 0 or w <= 0 or h <= 0:
            continue
        text = annotation['text']
        crop_img = image[y:y+h,x:x+w]
        crop_file_name = file_name[:-4]+'_{:03}.jpg'.format(idx+1)
        cv2.imwrite(save_root_path+'test/'+crop_file_name, crop_img)
        gt_file.write("test/{}\t{}\n".format(crop_file_name, text))

validation_annotations = json.load(open('./validation_annotation.json'))
gt_file = open(save_root_path+'gt_validation.txt', 'w')
for file_name in tqdm(validation_annotations):
    annotations = validation_annotations[file_name]
    image = cv2.imread(data_root_path+file_name)
    for idx, annotation in enumerate(annotations):
        x,y,w,h = annotation['bbox']
        if x<= 0 or y<= 0 or w <= 0 or h <= 0:
            continue        
        text = annotation['text']
        crop_img = image[y:y+h,x:x+w]
        crop_file_name = file_name[:-4]+'_{:03}.jpg'.format(idx+1)
        cv2.imwrite(save_root_path+'validation/'+crop_file_name, crop_img)
        gt_file.write("validation/{}\t{}\n".format(crop_file_name, text))
        
train_annotations = json.load(open('./train_annotation.json'))
gt_file = open(save_root_path+'gt_train.txt', 'w')
for file_name in tqdm(train_annotations):
    annotations = train_annotations[file_name]
    image = cv2.imread(data_root_path+file_name)
    for idx, annotation in enumerate(annotations):
        x,y,w,h = annotation['bbox']
        if x<= 0 or y<= 0 or w <= 0 or h <= 0:
            continue        
        text = annotation['text']
        crop_img = image[y:y+h,x:x+w]
        crop_file_name = file_name[:-4]+'_{:03}.jpg'.format(idx+1)
        cv2.imwrite(save_root_path+'train/'+crop_file_name, crop_img)
        gt_file.write("train/{}\t{}\n".format(crop_file_name, text))

 

다음 글에서 다룰 내용

이렇게 하고 나면 Text Recognition 모델을 학습시기키 위한 준비는 다 끝났다.

다음 글에서는 실제 학습을 진행하는 과정과 한글 학습을 위해 기존 모델 코드에서 변경된 점 그리고 학습 결과에 대한 내용을 다뤄보겠다.