[python] CRNN 한글 모델 학습하기(2) - Fine tuning 데이터셋

1. 최상위 폴더에 deep-text-recognition-benchmark 클론해오기
2. 최상위 폴더에 kor_dataset 폴더 생성 후, 데이터셋들을 해당 폴더로 옮기기
3. AIhub 데이터셋 학습하기
4. Finetuning 데이터셋 추가 학습하기
5. 완성된 pretrained model 테스트하기

AIhub에서 가져온 데이터셋이 너무 정직(?)해서, 좀 더 자연스러운 인쇄체+필기체들을 학습시키기 위해 추가로 데이터셋을 제작하였다.

이번 포스팅에서는 이전에 만든 kocrnn.pth 모델 위에 새로 제작한 데이터셋들을 추가로 학습시키는 과정을 정리한다.

현재 파일 구조는 아래와 같다.

최상위폴더
  |--deep-text-recognition-benchmark
    |--htr_data
      |--train
      |--test
      |--validation
      |--get_images.py
      |--gt_test.txt
      |--gt_train.txt
      |--gt_validation.txt
    |--ocr_data
      |--train
      |--test
      |--validation
      |--get_images.py
      |--gt_test.txt
      |--gt_train.txt
      |--gt_validation.txt
    |--...
  |--kor_dataset
    |--aihub_data
      |--htr
        |--images # 여기에 AIhub 필기체 이미지파일들 넣기
        |--handwriting_data_info1.json # 라벨링 파일
      |--ocr
        |--images # 여기에 AIhub 인쇄체 이미지파일들 넣기
        |--printed_data_info.json # 라벨링 파일
  |--saved_models
    |--kocrnn_ocr
      |--best_accuracy.pth
      ...
    |--kocrnn_ocr_htr
      |--best_accuracy.pth
      ...
  |--pretrained_models
    |--kocrnn.pth
  |--aihub_dataset.py
  |--...

데이터셋 가공하기

• 추가로 제작한 데이터셋을 아래 파일구조와 같이 넣어주도록 한다.

최상위폴더
  |--deep-text-recognition-benchmark
    |--htr_data
    |--ocr_data
    |--...
  |--kor_dataset
    |--aihub_data
    |--finetuning_data
      |--made1
        |--images      # 이미지파일들을 저장
        |--labels.txt  # 라벨링 파일
      |--made2
        |--images
        |--labels.txt
  |--saved_models
    |--kocrnn_ocr
      |--best_accuracy.pth
      ...
    |--kocrnn_ocr_htr
      |--best_accuracy.pth
      ...
  |--pretrained_models
    |--kocrnn.pth
  |--aihub_dataset.py
  |--...

• 참고로, 해당 데이터셋들의 라벨링 파일(labels.txt)의 구조는 아래와 같다.

# labels.txt
# {filename}.jpg {label}
0.jpg 여기가
1.jpg 라벨이
2.jpg 있는
3.jpg 곳입니다.
...

• 이전에 AIhub 데이터셋을 가공할 때와 마찬가지로 gt파일을 만들어주어야 한다. 최상위 폴더에 finetuning_dataset.py를 생성해준다.

# finetuning_dataset.py
import random
import os

data_type = 'made1' # made1, made2

# Data Preprocessing
with open(f'./kor_dataset/finetuning_data/{data_type}/labels.txt', 'r') as f:
  labels = f.readlines()

image_files = os.listdir('./kor_dataset/finetuning_data/{data_type}/images/')
total = len(image_files)

random.shuffle(image_files)

n_train = int(len(image_files) * 0.7)
n_validation = int(len(image_files) * 0.15)
n_test = int(len(image_files) * 0.15)

train_files = image_files[:n_train]
validation_files = image_files[n_train:n_train+n_validation]
test_files = image_files[-n_test:]

save_root_path = f'./deep-text-recognition-benchmark/{data_type}_data/'

gt_test = open(save_root_path+'gt_test.txt', 'w')
gt_validation = open(save_root_path + 'gt_validation.txt', 'w')
gt_train = open(save_root_path + 'gt_train.txt', 'w')

for line in labels:
    file_name, annotation = line.split('.jpg')
    file_name += '.jpg'
    print(file_name, annotation, end='')

    if file_name in train_files:
        gt_train.write("train/{}\t{}".format(file_name, annotation))
    if file_name in test_files:
        gt_test.write("test/{}\t{}".format(file_name, annotation))
    elif file_name in validation_files:
        gt_validation.write("validation/{}\t{}".format(file_name, annotation))

print('gt_file done')

• gt파일을 생성해준다.

python3 finetuning_dataset.py

• 이후 파일구조는 아래와 같다.

최상위폴더
  |--deep-text-recognition-benchmark
    |--htr_data
    |--ocr_data
    |--made1_data
      |--gt_test.txt
      |--gt_train.txt
      |--gt_validation.txt
    |--...
  |--kor_dataset
    |--aihub_data
    |--finetuning_data
      |--made1
        |--images      # 이미지파일들을 저장
        |--labels.txt  # 라벨링 파일
      |--made2
        |--images
        |--labels.txt
  |--saved_models
    |--kocrnn_ocr
      |--best_accuracy.pth
      ...
    |--kocrnn_ocr_htr
      |--best_accuracy.pth
      ...
  |--pretrained_models
    |--kocrnn.pth
  |--aihub_dataset.py
  |--finetuning_dataset.py
  |--...

• AIhub에서와 마찬가지로, deep-text-recognition-benchmark/made1_data에 get_images.py 파일을 만들고, train, test, validation 폴더를 생성한다.

cd deep-text-recognition-benchmark/made1_data
mkdir train test validation
# get_images.py 생성하기
python3 get_images.py

# get_images.py
import shutil

data_root_path = '../../kor_dataset/finetuning_data/made1/images/'
save_root_path = './images/'

# copy images from dataset directory to current directory
shutil.copytree(data_root_path, save_root_path)

# separate dataset : train, validation, test
obj_list = ['train', 'test', 'validation']
for obj in obj_list:
  with open(f'gt_{obj}.txt', 'r') as f:
    lines = f.readlines()
    for line in lines:
      file_path = line.split('.jpg')[0]
      file_name = file_path.split('/')[1] + '.jpg'
      res = shutil.move(save_root_path+file_name, f'./{obj}/')

• lmdb 데이터를 생성해준다.

cd ../../
python3 ./deep-text-recognition-benchmark/create_lmdb_dataset.py \
    --inputPath ./deep-text-recognition-benchmark/made1_data/ \
    --gtFile ./deep-text-recognition-benchmark/made1_data/gt_train.txt \
    --outputPath ./deep-text-recognition-benchmark/made1_data_lmdb/train

python3 ./deep-text-recognition-benchmark/create_lmdb_dataset.py \
    --inputPath ./deep-text-recognition-benchmark/made1_data/ \
    --gtFile ./deep-text-recognition-benchmark/made1_data/gt_validation.txt \
    --outputPath ./deep-text-recognition-benchmark/made1_data_lmdb/validation

이후 생성한 데이터셋들 역시 made#에서 #부분만 다른 숫자를 넣어 위와 같이 진행하면 된다.

추가 학습하기

• 학습 순서는 made1 > made2로 진행하였다.

추가학습 시작하기

# made1
CUDA_VISIBLE_DEVICES=0 python3 ./deep-text-recognition-benchmark/train.py \
    --train_data ./deep-text-recognition-benchmark/made1_data_lmdb/train \
    --valid_data ./deep-text-recognition-benchmark/made1_data_lmdb/validation \
    --Transformation TPS --FeatureExtraction ResNet --SequenceModeling BiLSTM --Prediction CTC \
    --batch_size 512 --batch_max_length 200 --data_filtering_off --workers 0 \
    --saved_model ./pretrained_models/kocrnn.pth --num_iter 100000 --valInterval 100

• 학습이 완료되면 여기에서 했던 것과 마찬가지로, saved_models에 저장된 폴더 명을 바꿔서 백업을 해두고, made2의 추가학습을 위해 best_accuracy.pth 모델을 따로 복사해둔다.
• made2까지 학습이 끝나면, 아래와 같은 파일구조가 된다.

최상위폴더
  |--deep-text-recognition-benchmark
    |--htr_data
    |--ocr_data
    |--made1_data
    |--made2_data
    |--lmdb 폴더들...
    |--기타 파일들...
  |--kor_dataset
    |--aihub_data
      |--htr
      |--ocr
   |--finetuning_data
      |--made1
      |--made2
  |--saved_models
    |--kocrnn_ocr
    |--kocrnn_ocr_htr
    |--kocrnn_ocr_htr_made1
    |--kocrnn_ocr_htr_made12
  |--pretrained_models
    |--kocrnn.pth
  |--aihub_dataset.py
  |--finetuning_dataset.py
  |--기타 파일들...

---