Tensorflow에서 OCR 모델을 만들어보자

Tetrapod·2024년 5월 25일

OCR tensorflow

OCR recognition 모델

목록 보기

1/2

이 글에서는 데이터셋 구축부터 OCR 모델 학습까지 다루어 본다.
OCR모델은 Text Detection과 Text Recognition으로 나뉜다.
여기서는 Recognition만 해보고자 한다.

OCR 이란?

광학 문자 인식(OCR) 이란 사람이 쓰거나 기계로 인쇄한 문자 이미지를 기계가 읽을 수 있는 데이터로 변환하는 과정

숫자 이미지 데이터 수집

아래와 같이 숫자가 포함된 이미지를 인터넷 상에서 500장 이상 수집하였다.

데이터 라벨링

라벨링 작업은 python의 'labelme'라는 툴을 이용한다.
이미지 속 보이는 숫자들은 전부 bound-box와 label을 넣어준다.
이 작업을 수집한 이미지 전부 반복해준다.

데이터 전처리

이전과정에서 이미지와 라벨을 일대일이 아닌 다대다로 만들었다.
이제 이미지상에서 bounding-box부분을 crop하고 데이터를 일대일로 만들어야한다.
crop 과정에서 이미지 시점 변환을 하였는데 OpenCV의 perspective 기능을 이용하였다.
대략 2천장의 데이터셋을 만들었다.
자세한 코드는 글 아래 링크 참조

이미지 증축

2천장의 이미지로 모델을 학습하기에 부족하므로 데이터 증축을 한다.
줌아웃, 회전, 기울임, 밝기, 이동, 노이즈 등을 무작위로 조절하여 이미지를 증축하였다.
자세한 코드는 글 아래 링크 참조

Loss 함수 정의

import tensorflow as tf

loss_fn = tf.keras.backend.ctc_batch_cost

def CTC_loss(y_true, y_pred):
    batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64") # 16
    input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64") # 50
    label_length = tf.cast(tf.shape(y_true)[1], dtype="int64") # 5

    input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
    label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")

    loss = loss_fn(y_true, y_pred, input_length, label_length)
    return loss

Tensorflow에는 ctc loss 함수가 내장되어있어 이를 사용한다.
이 손실 함수는 입력 시퀀스를 출력 시퀀스로 매핑하는 데 사용된다.
예를 들어, 음성을 입력받아 텍스트를 출력하거나, 이미지에서 텍스트를 인식하는 OCR 시스템에서는 CTC Loss를 사용할 수 있다.

모델 정의

def build_model(input_shape=(MAX_WIDTH, IMG_HEIGHT, 3)):
    input_img = layers.Input(shape=input_shape)
    
    arg_dic = {"activation":"swish", "padding":"same", "kernel_initializer":"he_normal"}
    x = layers.Conv2D(32, 3, **arg_dic)(input_img)
    x = layers.MaxPooling2D()(x)

    x = layers.Conv2D(64, 3, **arg_dic)(x)
    x = layers.MaxPooling2D()(x)

    arg_dic["activation"] = None
    x = layers.Conv2D(64, 3, **arg_dic)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("swish")(x)
    x = layers.Dropout(0.15)(x)
    x = layers.MaxPooling2D()(x)

    x = layers.Reshape(target_shape=(MAX_WIDTH // 8, -1), name="Reshape")(x)

    # x = layers.Conv1D(128, 1, **arg_dic)(x) 
    x = layers.Conv1D(256, 3, **arg_dic)(x) 
    x = layers.Conv1D(128, 1, activation=None, kernel_initializer="he_normal")(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("swish")(x)
    x = layers.Dropout(0.2)(x)

    # RNNs
    x = layers.Bidirectional(layers.LSTM(128, return_sequences=True, dropout=0.25))(x)
    x = layers.Bidirectional(layers.LSTM(64, return_sequences=True, dropout=0.25))(x)
    x = layers.Bidirectional(layers.LSTM(32, return_sequences=True, dropout=0.25))(x)

#     # Output layer
    # x = layers.Dense(len(char_to_num.get_vocabulary()) + 1,
    #                  activation="softmax")(x)
    x = layers.Conv1D(len(char_to_num.get_vocabulary()) + 1, 1, activation='softmax')(x)

    # Define the model
    model = tf.keras.models.Model(inputs=input_img, outputs=x, name="ocr_model_v1")
    return model