[DSBA]예비프로젝트 : OCR모델 후처리

fine-tunnig결과 성능이 저조한 커스텀 모델 대신 easyOCR 기본 모델을 사용하기로 했습니다.

기본 모델의 글자 인식 및 추출 능력은 우수했으나 한 가지 문제가 있었는데,

가로 줄 단위로 글자를 인식하다보니 어떤 단어가 줄 바꿈으로 떨어지면 각각 다른 단어가 되버렸습니다.

import easyocr
from PIL import Image, ImageDraw

img_path = './cosmetics_00026_crop.jpg'
reader = easyocr.Reader(['ko','en'], gpu=False)
detected_result = reader.detect(img_path)

# Load the image
img = Image.open(img_path)
# Create an ImageDraw object
draw = ImageDraw.Draw(img)
# Define the coordinates of the rectangle
for coords in detected_result[0][0]:
    # Draw the rectangle on the image
    coords = [coords[0], coords[2], coords[1], coords[3]]
    draw.rectangle(coords, outline='red', width=2)
img.save('./crop1_detect.jpg')

crop1_detect.jpg

기본 모델을 사용한 글자 추출

우리 팀은 화장품의 성분이 사용자의 피부타입과 니즈에 적합한지 아닌지 판단하는 시스템을 만들어야했기때문에 정확한 성분명이 추출되야했습니다. 

그래서 줄바꿈으로 분리된 성분명을 하나로 합쳐야했습니다. 위 사진만 보면 간단한 알고리즘으로 가능할 것 같지만,

다른 화장품 사진의 경우 정확한 글자가 추출되지 않은 경우도 있어서 단순히 단어를 이어붙이면 엉뚱한 성분명이 되버렸습니다.

cosmetics_18798.jpg

줄바꿈된 단어뿐만 아니라 다른 단어들도 정확히 추출되진 않아서 전체 단어를 정확한 성분명으로 바꾸는 알고리즘이 필요했습니다. 

팀원들과 함께 해결방법을 찾았는데, 'Levenshtein distance'라는 알고리즘을 알게되었습니다. Levenshtein distance는 단어의 편집거리를 구해 단어간 형태적 유사도를 측정하는 dynamic programming 알고리즘입니다.

이를 파이썬으로 구현한 라이브러리를 사용했습니다.

라이브러리 설치는 다음과 같습니다.

pip install python-Levenshtein

그리고 저희 클래스에 가장 유사도가 높은 단어를 찾는 함수를 다음과 같이 작성했습니다.

def findSimilar(self, words, user_dictionary):
    import numpy as np
    import Levenshtein
    temp = []
    for x in words:
        distances = np.array([Levenshtein.distance(x, w) for w in user_dictionary])
        try:
            index = np.argmin(distances)
            similar_word = user_dictionary[index]
            temp.append((x, similar_word, min(distances)))
        except ValueError:
            temp.append(w)    
    return temp

그리고 이를 이용해 분리된 단어를 합치는 함수를 만들었습니다.

# 분리된 단어 합치고 거리값 기반으로 수정하기
def concateUpdate(self, ocr_result, user_dictionary):
    line = []
    word_list = []
    for i in range(len(ocr_result)-1):
        line1 = ocr_result[i].split()
        line2 = ocr_result[i+1].split()
        if len(line1)*len(line2) != 0:
            connected_word = line1[-1]+line2[0]
            word = (line1[-1], line2[0], connected_word)
            word_list.append(word)
        if i == 0:
            line1, line2 = line1[:-1], line2[1:-1]
        elif i == len(ocr_result)-2:
            line1, line2 = line1[1:-1], line2[1:]
        else:
            line1, line2 = line1[1:-1], line2[1:-1]
        line.extend(line1+line2)
        line.append(connected_word)
        tot_sum = list(set(line))
    n = 0
    for wt in word_list:
        res = self.findSimilar(wt, user_dictionary)
        # ICDAR2019 Normalized Edit Distance 계산
        b1 = 1 - res[0][2] / max(len(res[0][0]), len(res[0][1]))
        b2 = 1 - res[1][2] / max(len(res[1][0]), len(res[1][1]))
        before_normED = (b1 + b2) / 2
        after_normED = 1 - res[2][2] / max(len(res[2][0]), len(res[2][1]))

        if before_normED >= after_normED:
            tot_sum.pop(tot_sum.index(res[2][0]))
            tot_sum.append(res[0][1])
            tot_sum.append(res[1][1])
            n += 1
        else:
            continue
    calculated_n = 2 * n
    return tot_sum, calculated_n

user_dictionary는 모든 화장품 성분이 있는 txt파일입니다. 대한화장품협회 사이트에서 성분사전 파일을 변환했습니다.

 

첫 번째 for문은 줄바꿈으로 떨어진 단어를 붙입니다.

이때 첫 줄의 첫 단어와 마지막줄의 마지막 단어는  따로 처리하고 그 외 줄바꿈 단어들은 떨어진 단어들 2개와 이어붙인 단어 1개, 총 3개가 word변수에 튜플 타입으로 저장됩니다.

 

이렇게 한 이유는 이어붙인 단어가 하나의 성분명이 아닌 경우도 존재하기 때문입니다. 그래서 떨어진 단어들과 합친 단어 각각 유사도를 측정해야합니다.

 

두 번재 for문에서 튜플 단어셋이 들어있는 리스트가 Levenshtein 거리를 계산하는 함수로 넘겨져 각 단어들마다 유사단어와 거리 값을 리턴 받습니다.

 

그리고 반환 받은 리스트는 ICDAR2019 normalized edit distance 계산 방식을 사용했습니다.

before_normED는 따로 떨어진 단어들의 normED 평균값입니다.

after_normED는 이어붙인 단어의 normED입니다.

 

그 아래 if문에서는 before_normED가 after_normED 보다 작거나 같다면 이어붙인 단어를 기존 단어와 이어붙인 단어를 모아놓은 리스트인 tot_sum에서 제거하고 떨어진 단어 2개를 추가합니다.

 

n은 이어붙인 단어가 다시 떨어질때마다 횟수를 셉니다. 그 이유는 줄바꿈 단어가 아닌 다른 단어들도 나중에 성분사전에서 유사도가 높은 단어로 바꿀때 사용하기 위함입니다. if문에서 추가한 떨어진 단어는 이미 성분사전에서 유사한 단어로 바뀐 단어이기 때문에 n을 이용해 리스트 슬라이싱을 하면 이를 다시 계산하지 않을 수 있습니다.

 

아래는 OCR후처리 전체 코드입니다.

class OCR:
    def __init__(self):
        import easyocr
        import torch
        if torch.cuda.is_available():
            self.reader = easyocr.Reader(['ko', 'en'], gpu=True)
        else:
            self.reader = easyocr.Reader(['ko', 'en'])
    
    def findSimilar(self, words, user_dictionary):
        import numpy as np
        import Levenshtein
        temp = []
        for x in words:
            distances = np.array([Levenshtein.distance(x, w) for w in user_dictionary])
            try:
                index = np.argmin(distances)
                similar_word = user_dictionary[index]
                temp.append((x, similar_word, min(distances)))
            except ValueError:
                temp.append(w)    
        return temp
    
    # 분리된 단어 합치고 거리값 기반으로 수정하기
    def concateUpdate(self, ocr_result, user_dictionary):
        line = []
        word_list = []
        for i in range(len(ocr_result)-1):
            line1 = ocr_result[i].split()
            line2 = ocr_result[i+1].split()
            if len(line1)*len(line2) != 0:
                connected_word = line1[-1]+line2[0]
                word = (line1[-1], line2[0], connected_word)
                word_list.append(word)
            if i == 0:
                line1, line2 = line1[:-1], line2[1:-1]
            elif i == len(ocr_result)-2:
                line1, line2 = line1[1:-1], line2[1:]
            else:
                line1, line2 = line1[1:-1], line2[1:-1]
            line.extend(line1+line2)
            line.append(connected_word)
            tot_sum = list(set(line))
        n = 0
        for wt in word_list:
            res = self.findSimilar(wt, user_dictionary)
            # ICDAR2019 Normalized Edit Distance 계산
            b1 = 1 - res[0][2] / max(len(res[0][0]), len(res[0][1]))
            b2 = 1 - res[1][2] / max(len(res[1][0]), len(res[1][1]))
            before_normED = (b1 + b2) / 2
            after_normED = 1 - res[2][2] / max(len(res[2][0]), len(res[2][1]))
            
            if before_normED >= after_normED:
                tot_sum.pop(tot_sum.index(res[2][0]))
                tot_sum.append(res[0][1])
                tot_sum.append(res[1][1])
                n += 1
            else:
                continue
        calculated_n = 2 * n
        return tot_sum, calculated_n
    
    def ocrWord(self, image, user_dictionary):
        result = self.reader.readtext(image, detail=0)
        updated_result, calculated_n = self.concateUpdate(result, user_dictionary)
        pass_result = updated_result[:-calculated_n]
        if calculated_n == 0:
            pass_result = updated_result
        changed_result = [x[1] for x in self.findSimilar(pass_result, user_dictionary)]
        final_result = list(set(changed_result + updated_result[-calculated_n:]))
        return final_result

이 클래스의 사용법은 다음과 같습니다.

# txt 파일 불러오기
with open("./user_dictionary.txt", "r", encoding='utf-8') as file:
    strings = file.readlines()
user_dictionary = [x.replace('\n','') for x in strings]

img_path = './cosmetics_18798.jpg'

ocr = OCR()
igd = ocr.ocrWord(img_path, user_dictionary)
print(igd)

기본 모델만으로 정확히 추출되지 않았던 성분명들이 이제는 성분사전에 있는 정확한 성분명들로 바뀐 것을 볼 수 있습니다. 

 

제가 만든 이 알고리즘은 기본모델이 어느정도는 알아볼수 있게 문자를 추출해야 신뢰도가 높은 결과를 얻을 수 있습니다.

그리고 성분 사전에서 유사단어를 찾는 코드가 완전탐색이어서 그리 빠른 편은 아닙니다.

이 알고리즘을 만들 당시에는 알고리즘에 대한 지식이 거의 없어서 더 빠르고 효율적인 방법을 생각해내기 어려웠습니다.

그래도 며칠간 고민하던 문제를 해결했을때 느꼈던 성취감과 보람이 있었기에 프로젝트를 성공적으로 마칠 수 있었습니다.

 

이상으로 예비프로젝트 포스팅을 마치겠습니다.