[ OCR ] 문자 추출 및 인식 (EAST text Detector Model) - Python

일반적으로 OCR에서 원하는 영역을 추출하는 방법은 이전 글에서 설명드린 것처럼, 그래디언트를 적용하고 단락으로 그룹화하여 조건식으로 특징을 입력하여 찾습니다. 

문제는 통제되지 않은 촬영 환경에서 촬영한 사진에서 텍스트를 감지하는 것은 매우 어렵다는 것입니다. 시야각은 텍스트와 평행하지 않고 빛이나 조명에 의해 지나치게 어둡거나 밝을 수도 있습니다.

EAST 문자 감지 모델은 Text가 흐리거나 기울어진 환경에서도, 부분적으로 가려직 경우에도 찾아낼 수 있습니다. EAST는 An Efficient and Accurate Scene Text Detector의 약자로 Zhou et al.의 2017년 논문을 기반으로 하는 딥러닝 문자 감지기입니다.

논문에 따르면 EAST는 720p 이미지에서 방향과 상관없이 인식할 수 있으며 13.2fps로 실행 할 수 있습니다. 속도가 빠르기때문에 정지 영상뿐 아니라 동영상에서도 사용이 가능합니다.

---

시작하기에 앞서 EAST 모델 다운로드가 필요합니다.

Import Packages

from imutils.object_detection import non_max_suppression
from imutils.perspective import four_point_transform
from imutils.contours import sort_contours
import matplotlib.pyplot as plt
import imutils
import numpy as np
import requests
import pytesseract
import cv2

Function

Colab에서 이미지를 확인하기위한 Function과 추출된 결과를 가공하는 Function입니다.

def plt_imshow(title='image', img=None, figsize=(8 ,5)):
    plt.figure(figsize=figsize)

    if type(img) == list:
        if type(title) == list:
            titles = title
        else:
            titles = []

            for i in range(len(img)):
                titles.append(title)

        for i in range(len(img)):
            if len(img[i].shape) <= 2:
                rgbImg = cv2.cvtColor(img[i], cv2.COLOR_GRAY2RGB)
            else:
                rgbImg = cv2.cvtColor(img[i], cv2.COLOR_BGR2RGB)

            plt.subplot(1, len(img), i + 1), plt.imshow(rgbImg)
            plt.title(titles[i])
            plt.xticks([]), plt.yticks([])

        plt.show()
    else:
        if len(img.shape) < 3:
            rgbImg = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        else:
            rgbImg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        plt.imshow(rgbImg)
        plt.title(title)
        plt.xticks([]), plt.yticks([])
        plt.show()

def decode_predictions(scores, geometry):
	(numRows, numCols) = scores.shape[2:4]
	rects = []
	confidences = []

	for y in range(0, numRows):
		scoresData = scores[0, 0, y]
		xData0 = geometry[0, 0, y]
		xData1 = geometry[0, 1, y]
		xData2 = geometry[0, 2, y]
		xData3 = geometry[0, 3, y]
		anglesData = geometry[0, 4, y]

		for x in range(0, numCols):
			if scoresData[x] < min_confidence:
				continue

			(offsetX, offsetY) = (x * 4.0, y * 4.0)

			angle = anglesData[x]
			cos = np.cos(angle)
			sin = np.sin(angle)

			h = xData0[x] + xData2[x]
			w = xData1[x] + xData3[x]

			endX = int(offsetX + (cos * xData1[x]) + (sin * xData2[x]))
			endY = int(offsetY - (sin * xData1[x]) + (cos * xData2[x]))
			startX = int(endX - w)
			startY = int(endY - h)

			rects.append((startX, startY, endX, endY))
			confidences.append(scoresData[x])

	return (rects, confidences)

Load Model

layerNames = ["feature_fusion/Conv_7/Sigmoid",
              "feature_fusion/concat_3"]

# 사전에 훈련된 EAST text detector 모델 Load
print("[INFO] loading EAST text detector...")
net = cv2.dnn.readNet("model/frozen_east_text_detection.pb")

width = 640
height = 640
min_confidence = 0.5
padding = 0.0

Load Image

url = 'https://user-images.githubusercontent.com/69428232/149087561-4803b3e0-bcb4-4f9f-a597-c362db24ff9c.jpg'

image_nparray = np.asarray(bytearray(requests.get(url).content), dtype=np.uint8)
org_image = cv2.imdecode(image_nparray, cv2.IMREAD_COLOR) 

plt_imshow("Original", org_image)

Image Resize

orig = org_image.copy()
(origH, origW) = org_image.shape[:2]

(newW, newH) = (width, height)
rW = origW / float(newW)
rH = origH / float(newH)

org_image = cv2.resize(org_image, (newW, newH))
(H, W) = org_image.shape[:2]

Text Detection

blob = cv2.dnn.blobFromImage(org_image, 1.0, (W, H), (123.68, 116.78, 103.94), swapRB=True, crop=False)
net.setInput(blob)
(scores, geometry) = net.forward(layerNames)

(rects, confidences) = decode_predictions(scores, geometry)
boxes = non_max_suppression(np.array(rects), probs=confidences)

Recognition

results = []

for (startX, startY, endX, endY) in boxes:
	startX = int(startX * rW)
	startY = int(startY * rH)
	endX = int(endX * rW)
	endY = int(endY * rH)

	dX = int((endX - startX) * padding)
	dY = int((endY - startY) * padding)

	startX = max(0, startX - dX)
	startY = max(0, startY - dY)
	endX = min(origW, endX + (dX * 2))
	endY = min(origH, endY + (dY * 2))

	# 영역 추출
	roi = orig[startY:endY, startX:endX]

	config = ("-l eng --psm 4")
	text = pytesseract.image_to_string(roi, config=config)

	results.append(((startX, startY, endX, endY), text))

Result

results = sorted(results, key=lambda r:r[0][1])

output = orig.copy()

# 결과 출력
for ((startX, startY, endX, endY), text) in results:
	cv2.rectangle(output, (startX, startY), (endX, endY), (0, 255, 0), 5)
	cv2.putText(output, text, (startX, startY - 20), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 5)


plt_imshow("Text Detection", output, figsize=(16, 10))

결과는 보시는 것과 같이 어느정도 잘 찾은 걸로 보입니다. 인식은 전처리 없이 Tesseract를 이용했으니 정확도는 떨어지네요.

Source Code

[ Colab ] EAST text detector

---

영상에 적용해도 무리없이 수행됩니다.