[Python] 사과게임 이미지에서 숫자 추출하기

사과게임을 하다가 모든 사과를 다 제거할 수 있는 맵이 존재하는지 궁금해서

파이썬으로 해당 맵을 체크할 수 있는지 확인하는 프로그램을 작성하려고 한다.

개발환경

Colab

구글 드라이브 연동

from google.colab import drive
drive.mount('/content/drive')

 

개발환경은 간단하게 Colab을 이용했으며, 환경 구성을 위해 구글 드라이브를 연동하고 OpenCV와 Tesseract 등을 추가적으로 설치했다.

OpenCV 및 Tesseract 설치

!pip install opencv-python
!sudo apt install tesseract-ocr
!pip install pytesseract

 

OpenCV는 실시간 컴퓨터 비전을 목적으로 하는 오픈소스 라이브러리이다. C/C++ 언어로 개발되었으며 파이썬, 자바 등에 바인딩되어 다양한 환경에서 사용 가능하다. 사과게임에서 숫자 이미지를 텍스트로 변환하기 전에 인식이 잘 될수 있도록 전처리를 하기위해 사용했다.

Tesseract는 다양한 운영체제에서 사용할 수 있는 광학 문자 인식 엔진이다. C++ 언어로 작성되었으며, 파이썬 환경에서 사용할 수 있도록 pytersseract 라이브러리를 설치했다. 숫자 이미지를 텍스트로 변환하기 위해 사용했지만 인식률이 좋지 않아서 포기했다. 추후에 언어 데이터 혹은 모델을 학습시키는 등 인식률을 높이는 방법들을 알아보면 좋을 것 같다.

 

WARNING: The following packages were previously imported in this runtime:  
  [PIL]  
You must restart the runtime in order to use newly installed versions.

-> RESTART RUNTIME

 

Colab 환경에서 OpenCV와 Tesseract를 설치하면 다음과 같은 경고메시지가 나오면서 제대로 설치가 되지 않는데, RESTART RUNTIME을 클릭하여 런타임을 재실행하면 된다.

라이브러리

import cv2
from google.colab.patches import cv2_imshow
import pytesseract
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import MinMaxScaler

print(f"cv2 {cv2.__version__}")
print(f"pytesseract {pytesseract.__version__}")

cv2 4.6.0
pytesseract 0.3.10

 

Colab 환경에서 imshow() 함수를 사용하면 다음과 같은 에러가 발생한다.

 

DisabledFunctionError: cv2.imshow() is disabled in Colab, because it causes Jupyter sessions
to crash; see https://github.com/jupyter/notebook/issues/3935.
As a substitution, consider using
  from google.colab.patches import cv2_imshow

 

Jupyter Notebook에서 imshow()를 사용하면 Jupyter session이 충돌되는 이슈로 Colab에서는 이를 대신하여 google.colab.patches에서 cv2_imshow()를 제공한다.

Colab에서 제공하는 클라우드 환경이 아니라 로컬 Jupyter Notebook 환경이라면 matplotlib.pyplot에서 제공하는 imshow()를 사용하거나 cv2_imshow()에서 PIL 라이브러리를 활용하는 코드를 활용하면 될 것 같다.
그 외에도 cv2.waitKey(0)와 cv2.destroyAllWindows()을 활용하여 해결하는 방법도 있다.
➡ How to use OpenCV imshow() in a Jupyter Notebook — Quick Tip

이미지에서 숫자 추출하기

1. pytesseract를 이용한 문자인식

1.1. 이미지 가져오기

img_path = "/content/drive/MyDrive/apple/apple.png"
img_cv = cv2.imread(img_path)
cv2_imshow(img_cv)

imread(const string & filename, int flags)

• filename : Name of file to be loaded
• flags : Flag that can take values of cv::ImreadModes |

OpenCV에서는 이미지를 BGR 형식의 numpy 배열로 이미지를 저장한다. matplotlib 등 일반적인 색배열은 RGB를 사용하기 때문에, 이를 일반적으로 출력하면 Blue와 Red가 뒤바뀌게 된다.

 

from IPython import display
import PIL
display.display(PIL.Image.fromarray(img_cv))

 

array([[[239, 254, 241], 
        [239, 254, 241],
        [239, 254, 241],
        ...,
        [239, 254, 241],
        [239, 254, 241],
        [239, 254, 241]],
       ...,
       [[208, 245, 219],
        [208, 245, 219],
        [208, 245, 219],
        ...,
        [216, 247, 224],
        [232, 251, 236],
        [239, 254, 241]]], dtype=uint8)

 

따라서 matplotlib 등 다른 라이브러리를 사용하기 위해서는 BGR 배열을 RGB로 바꿔줘야한다. google.colab.patches에서 제공하는 cv2_imshow() 또한 pillow(PIL)를 사용하기 때문에 BGR을 RGB로 변환하여 이미지를 출력한다.

 

from IPython import display
import PIL
img_rgb = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
display.display(PIL.Image.fromarray(img_rgb))

 

array([[[241, 254, 239],
        [241, 254, 239],
        [241, 254, 239],
        ...,
        [241, 254, 239],
        [241, 254, 239],
        [241, 254, 239]],
       ...,
       [[219, 245, 208],
        [219, 245, 208],
        [219, 245, 208],
        ...,
        [224, 247, 216],
        [236, 251, 232],
        [241, 254, 239]]], dtype=uint8)

 

1.2. 이미지 흑백 및 블러처리

노이즈를 줄이고 연산속도를 향상시키기 위해서 이미지를 흑백 및 블러처리.
imread()의 flag를 IMREAD_GRAYSCALE 혹은 0으로 설정하는 방법도 있다.

# img_path = "/content/drive/MyDrive/apple/apple.png"
# img_gray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
img_gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
img_blurred = cv2.GaussianBlur(img_gray, (5,5), 0)
cv2_imshow(img_blurred)

 

1.3. 이미지 임계처리

img_adaptiveThreshold = cv2.adaptiveThreshold(
    img_blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 61, -60
)
cv2_imshow(img_adaptiveThreshold)

 

1.4. 이미지 팽창 및 침식

# img_morphologyEx = cv2.morphologyEx(img_adaptiveThreshold, cv2.MORPH_CLOSE, k)
# cv2_imshow(img_morphologyEx)

k = cv2.getStructuringElement(cv2.MORPH_RECT, (1,1))
k2 = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
img_dilate = cv2.dilate(img_adaptiveThreshold, k)
img_erode = cv2.erode(img_dilate, k2)
cv2_imshow(img_erode)

 

1.5. 텍스트 추출

text = pytesseract.image_to_string(img_morphologyEx, config="--psm 6").replace(" ", "")
print(text)

 

66292691323834318
89358479535813156
22575945384595889
7438266572552531«4
65722751196488534
52149252498553672
41583252173141964
73471966487528784
54918791691639861
26287444455448728

 

4번째 줄을 보면 숫자 이미지가 노이즈때문에 제대로 인식이 안되어서 숫자 추출이 실패했다.

노이즈를 줄일 수 있는 방법을 찾아보다가 어차피 숫자 이미지가 정해져있다는 생각이 들어서

차라리 숫자 이미지를 사용해서 숫자를 추출하는 방식이 좋을 것 같다고 생각했다.

2. 추출된 숫자 이미지로 탬플릿 매칭

2.1. 이미지 엣지검출

img_canny = cv2.Canny(img_gray, 100, 200)
cv2_imshow(img_canny)

2.2. 이미지 윤곽처리

cnts, hierarchy = cv2.findContours(img_canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

img_rect = img_cv.copy()

array=[]
for i in range(len(cnts)):
    cnt = cv2.boundingRect(cnts[i])
    x,y,w,h = cnt
    rect_area = w*h
    aspect_ratio = (float)(w)/h

    if(aspect_ratio<0.8) and (aspect_ratio>0.2) and (rect_area > 300) and (rect_area < 600):
        array.append(cnt)
        cv2.rectangle(img_rect, (x,y), (x+w,y+h), (255,0,0), 2)

array = set(array)
print(f"array {len(array)}")
if(len(array)!=170):
    raise Exception("중복 인식된 사과가 있습니다.")
cv2_imshow(img_rect)

 

2.3. 숫자 이미지 추출

array=list(array)
i=23 # 1-9까지 하나씩 해야함...
img_crop = img_cv[array[i][1]:array[i][1]+array[i][3],array[i][0]:array[i][0]+array[i][2]]
cv2_imshow(img_crop)

2.4. 추출된 숫자 이미지로 좌표 데이터 확인

array=[]
for i in range(9,0,-1):
    one_path = f"/content/drive/MyDrive/apple/사과/{i}.jpg"
    tmplt = cv2.imread(one_path,0)
    w, h = tmplt.shape[::-1]

    result = cv2.matchTemplate(img_gray, tmplt, cv2.TM_CCOEFF_NORMED)
    threshold = 0.99 # 임계치 설정
    box_loc = np.where(result >= threshold) # 임계치 이상의 값들만 사용

    img_rect = img_cv.copy()

    for box in zip(*box_loc[::-1]):
        startX, startY = box
        endX, endY = startX + w, startY + h
        array.append([i,float(startX),float(startY)])
        cv2.rectangle(img_rect, (startX, startY), (endX, endY), (255,0,0), 2)

if(len(array)!=170):
    raise Exception("중복 인식된 사과가 있습니다.")

cv2_imshow(img_rect)

data=np.array(array)
plt.scatter(data[:,-2], data[:,-1])
plt.show()

 

2.5. 좌표 데이터 정규화

df = pd.DataFrame(data)
df.columns = ["weight", "x", "y"]

scalerX = MinMaxScaler((0,16))
scalerY = MinMaxScaler((0,9))

df['x'] = scalerX.fit_transform(df["x"].values.reshape(-1,1))
df['y'] = scalerY.fit_transform(df["y"].values.reshape(-1,1))

df = df.round(0).astype(int)

plt.xlim([-1, 17])      # X축의 범위: [xmin, xmax]
plt.ylim([-1, 10])     # Y축의 범위: [ymin, ymax]
plt.xticks([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17])
plt.yticks([0,1,2,3,4,5,6,7,8,9,10])

plt.scatter(df['x'], df['y'])
plt.show()
df.head()

 

2.6. 좌표 데이터 2차원 배열

apple = np.zeros((10,17), dtype=tuple)
for i in range(len(df)):
    n,x,y = df.loc[i]
    apple[y][x] = (i, n)

for i in apple:
    for j in i:
        print(j[1], end=" ")
    print("")

 

6 6 2 9 2 6 9 1 3 2 3 8 3 4 3 1 8 
8 9 3 5 8 4 7 9 5 3 5 8 1 3 1 5 6 
2 2 5 7 5 9 4 5 3 8 4 5 9 5 8 8 9 
7 4 3 8 2 6 6 5 7 2 5 5 2 5 3 1 4 
6 5 7 2 2 7 5 1 1 9 6 4 8 8 5 3 4 
5 2 1 4 9 2 5 2 4 9 8 5 5 3 6 7 2 
4 1 5 8 3 2 5 2 1 7 3 1 4 1 9 6 4 
7 3 4 7 1 9 6 6 4 8 7 5 2 8 7 8 4 
5 4 9 1 8 7 9 1 6 9 1 6 3 9 8 6 1 
2 6 2 8 7 4 4 4 4 5 5 4 4 8 7 2 8