[졸업프로젝트] 상품 태그 OCR 모델 개발 과정

이번 졸업프로젝트에서 AI 모듈 구현을 맡게 되어서, OCR을 활용한 상품 태그 인식 파이프라인을 구현했다.

우리 팀이 개발하는 서비스 'TruthTag'는 식료품 스마트컨슈밍 어시스턴트로, 가공 식료품에 스킴플레이션이나 슈링크플레이션이 일어나는지를 모니터링하고, 상품의 변동 사항을 시각화하여 제공하는 서비스이다. 사용자가 원한다면 특정 상품 대신 구매할 수 있는 대체품을 추천하는 챗봇 서비스 기능도 있다. 그러니까 우리 서비스에서 AI 모듈이 필요한 곳은 1) 사용자가 이미지로 상품 검색과 2) 챗봇 두 군데인데, 이번 게시글에서는 1) 이미지로 상품 검색 기능을 개발한 과정에 대해 작성해보려고 한다.

0. 기능 설명

사용자가 마트나 편의점에서 상품 태그 사진을 찍으면, OCR을 통해 이 사진 내의 텍스트를 인식하여 딕셔너리 형태로 반환하는 기능을 구현했다.

1. 데이터 수집

데이터는 직접 마트와 편의점을 돌아다니면서 촬영한 상품 태그 이미지 133장이다. 판매처마다 상품 태그의 구성이 달랐기 때문에, 모델의 강건성을 위해 4개의 다른 곳의 데이터를 수집했다.

2. 상품 태그 Segmentation

식료품점에서 바로 상품 사진을 찍으면 뒷배경에 있는 문자까지 인식되기 때문에, 인식하고자 하는 영역을 추출하는 과정이 필요하다. 이를 위해서 상품 태그 영역을 labeling하였고, YOLOv5 모델을 사용해 상품 태그 영역 인식 모델을 학습했다.

전처리로는 100장의 train data에 resizing, rotation, random noise 추가, contrast 조절 등의 augmentation 과정을 거친 뒤 학습에 활용했다.

200 epoch 동안 학습을 진행했고, 학습된 모델을 test한 결과, 다음과 같이 상품 태그 영역을 잘 segment한 것을 확인할 수 있다.

상품 태그 영역 인식이 끝나면, 인식한 영역의 좌표를 출력으로 얻게 된다. 이제 여기서 얻은 좌표를 활용해 OCR에 사용할 영역만을 crop한다. crop_img 변수에 이 좌표들을 사용해 crop한 이미지가 저장된다.

import torch
import cv2
import os

new_img = cv2.imread(img_path) # img_path에 OCR하고자 하는 이미지 경로 지정

# Run yolov5
model = torch.hub.load('ultralytics/yolov5', 'custom', path='../weights/best_200epoch.pt') # path 경로에 위에서 학습한 모델 경로 지정
results = model(new_img)

# save cropped area
results.crop()
cur_dir = os.getcwd()
crop_path = os.path.join(cur_dir, 'runs/detect/exp/crops/words/image0.jpg') # crop한 이미지 저장할 경로 지정
crop_img = cv2.imread(crop_path) # crop 이미지 load

3. 상품 태그 OCR

이제 OCR을 수행하면 된다. 이를 위해 EasyOCR 모델을 전이학습했는데, 자세한 학습 방법은 항목 4. EasyOCR 학습하기에서 설명하도록 하겠다. EasyOCR API로 글자 영역을 인식하고, 영역 내 글자들을 인식한다. lang_list 인자로 인식하고자 하는 언어의 종류를 입력하면 된다. 나는 한글과 영어를 인식하고자 했으므로 'ko', 'en'을 선택했다.

from easyocr import Reader

# run easyOCR in cropped
langs = ['ko', 'en']
print("[INFO] OCR'ing input image...")
reader = Reader(lang_list=langs, gpu=True)
results = reader.readtext(crop_img)

# loop over the results
real = []
for (bbox, text, prob) in results:
	if ')' in text or '>' in text or '"\"' in text or '원' in text:
        real.append(text)
	(tl, tr, br, bl) = bbox
	tl = (int(tl[0]), int(tl[1]))
	tr = (int(tr[0]), int(tr[1]))
	br = (int(br[0]), int(br[1]))
	bl = (int(bl[0]), int(bl[1]))

4. EasyOCR 학습하기

추가 학습을 하지 않은 EasyOCR 모델은 성능이 그다지 좋지 않기 때문에, 인식 성능을 높이기 위해 상품 태그 데이터를 라벨링하고 학습했다.

학습을 위해서는 CLOVAai의 https://github.com/clovaai/deep-text-recognition-benchmark repository를 참고했다. 모델은 VGG Feature Extractor와 Bidirectional LSTM을 선택했다.

학습하기 위해서는 왼쪽과 같이 이미지 주소와 text 데이터를 tab 기준으로 띄워서 저장한 후, 경로에 맞게 각 데이터를 위치시키면 된다.

Epoch 110 기준 accuracy는 3%, Train loss는 1.75 정도였으나, Epoch 1000 이상 학습 하자 batch 내의 accuracy가 100%, Train loss 0.00076를 달성했다.

5. 인식 결과 후처리

결과적으로 다음과 같은 형식의 문자 인식 결과를 서버에 넘겨줘야 했으므로, 인식 결과를 formatting하는 작업이 필요하다. 그래서 for문을 돌면서 result 변수에 담긴 OCR 결과들이 어떤 항목에 해당하는지 확인하고, 딕셔너리에 담아 반환했다.

{'company': '회사명', 'product': '상품명', 'price': '가격'}

info = dict()
company = ''
product = ''
price = ''
for i in results:
	if ')' in i:
		company = i.split(')')[0]
		product = i.split(')')[1]
		info['company'] = company
		info['product'] = product
	if '>' in i:
		price = i.split('>')[-1]
		info['price'] = price
	if "'\'" in i:
		price = i.split("'\'")[-1]
		info['price'] = price
	if '#' in i:
		price = i.split("#")[-1]
		info['price'] = price
	if "원" in i:
		price = i.split("원")[0]
		info['price'] = price
	if "," in i:
		price = str(i)
		info['price'] = price
	if "00" in i:
		price = str(i)
		info['price'] = price

이후 인식한 글자 중 '원' 단위를 나타내는 문자와 제조사 이름 뒤에 오는 ')' 등의 특수문자를 제거하는 과정을 거친다.

# STEP 1: 특수문자 제거
for i in info:
	info[i] = re.sub(r"[^\uAC00-\uD7A30-9a-zA-Z\s]", '', info[i])

EasyOCR의 글자 인식 결과가 매우 좋지는 않았기 때문에, 일종의 오타 수정 과정으로서 Fuzzy String Matching을 사용했다. 상품 제조사의 경우 종류가 많지 않으므로, 제조사 후보들을 미리 리스트에 담아둘 수 있었고, 인식 결과를 이 후보들 중 임베딩이 가장 비슷한 것으로 치환하는 과정을 거쳤다.

from rapidfuzz import fuzz

company_list = ["매일", "농심", "빙그레", "서울우유", "롯데", "오리온", "오뚜기", "리츠", "팔도", "스타벅스"]

# STEP 2: Fuzzy String Matching
idxi, idxj = -1, -1
min_company, min_product = 0, 0
for i in range(len(company_list)):
	prob = fuzz.ratio(company_list[i], info['company'])
	if prob > min_company:
		idxi = i
        
info['company'] = company_list[idxi]

6. 출력 예시

Test dataset 중 하나를 가지고 앞서 설명한 과정을 진행한 결과는 다음과 같다.

HK's log