import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"

import pickle

import glob

import string

import time

import math

import argparse

from shutil import copyfile, rmtree

import tqdm

import numpy as np

from numpy.random import RandomState

"""

###############################################################################################

# REFERENCES #

###############################################################################################

https://github.com/meijieru/crnn.pytorch

https://github.com/sbillburg/CRNN-with-STN/blob/master/CRNN_with_STN.py

https://github.com/keras-team/keras/blob/master/examples/image_ocr.py

https://github.com/keras-team/keras-applications/blob/master/keras_applications/mobilenet.py

CTC:

https://distill.pub/2017/ctc/

https://towardsdatascience.com/intuitively-understanding-connectionist-temporal-classification-3797e43a86c

Spatial transformer network:

https://github.com/oarriaga/STN.keras

https://arxiv.org/pdf/1506.02025.pdf

https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html

###############################################################################################

########

# RUN: #

########

docker build -t crnn_ocr:latest -f Dockerfile .

nvidia-docker run --rm -it -v /home/data/OCR/data/mjsynth/mnt/ramdisk/max/90kDICT32px:/input_data \

-v /home/data/OCR/data:/save_path \

-p 8000:8000 gasparjan/crnn_ocr:latest

nvidia-docker run --rm -it -v /home:/data \

-p 8004:8000 gasparjan/crnn_ocr:latest

____________

Mjsynth

____________

python3 train.py --G 1 --path /data/data/OCR/data/mjsynth/mnt/ramdisk/max/90kDICT32px --training_fname annotation_train.txt \

--val_fname annotation_test.txt --save_path /data/data/OCR/data --model_name OCR_mjsynth_FULL --nbepochs 1 \

--norm --mjsynth --opt adam --time_dense_size 128 --lr .0001 --batch_size 64 --early_stopping 5000

python3 train.py --G 1 --path /data/data/OCR/data/mjsynth/mnt/ramdisk/max/90kDICT32px --training_fname annotation_train.txt \

--val_fname annotation_test.txt --save_path /data/data/OCR/data --model_name OCR_mjsynth_FULL_2 --nbepochs 1 \

--norm --mjsynth --opt adam --time_dense_size 128 --lr .0001 --batch_size 64 --early_stopping 20 \

--pretrained_path /data/data/OCR/data/OCR_mjsynth_FULL/checkpoint_weights.h5

____________

IAM

____________

python3 train.py --G 1 --path /data/data/CRNN_OCR_keras/data/IAM_processed --train_portion 0.9 --save_path ./data \

--model_name OCR_IAM_ver1 --nbepochs 200 --norm --opt adam --time_dense_size 128 --lr .0001 \

--batch_size 64 --pretrained_path /data/data/OCR/data/OCR_mjsynth_FULL_2/final_weights.h5

"""

if __name__ == '__main__':

parser = argparse.ArgumentParser(description='crnn_ctc_loss')

parser.add_argument('-p', '--path', type=str, required=True)

parser.add_argument('--training_fname', type=str, required=False, default=None)

parser.add_argument('--val_fname', type=str, required=False, default="")

parser.add_argument('--save_path', type=str, required=True)

parser.add_argument('--model_name', type=str, required=True)

parser.add_argument('--pretrained_path', default=None, type=str, required=False)

parser.add_argument('--nbepochs', type=int, default=20)

parser.add_argument('--G', type=str, default="1")

parser.add_argument('--random_state', type=int, default=42)

parser.add_argument('--train_portion', type=float, default=0.9)

parser.add_argument('--time_dense_size', type=int, default=128)

parser.add_argument('--n_units', type=int, default=256)

parser.add_argument('--batch_size', type=int, default=64)

parser.add_argument('--opt', type=str, default="sgd")

parser.add_argument('--lr', type=float, default=0.001)

parser.add_argument('--early_stopping', type=int, default=0)

parser.add_argument('--norm', action='store_true')

parser.add_argument('--mjsynth', action='store_true')

parser.add_argument('--GRU', action='store_true')

# default values set according to mjsynth dataset rules

parser.add_argument('--imgh', type=int, default=100)

parser.add_argument('--imgW', type=int, default=32)

args = parser.parse_args()

globals().update(vars(args))

os.environ["CUDA_VISIBLE_DEVICES"] = G

import tensorflow as tf

from keras import backend as K

from keras.callbacks import ModelCheckpoint, LearningRateScheduler

from keras.utils.training_utils import multi_gpu_model

from keras.models import load_model, clone_model

from keras.layers import Lambda

from utils import *

try:

rmtree(save_path+"/"+model_name)

except:

pass

os.mkdir(save_path+"/"+model_name)

with open(save_path+'/'+model_name+"/arguments.txt", "w") as f:

f.write(str(args))

prng = RandomState(random_state)

lexicon = get_lexicon()

classes = {j:i for i, j in enumerate(lexicon)}

inverse_classes = {v:k for k, v in classes.items()}

print(" [INFO] %s" % classes)

if mjsynth:

train = open(os.path.join(path, training_fname), "r").readlines()

train = parse_mjsynth(path, train)

prng.shuffle(train)

val = np.array(open(os.path.join(path, val_fname), "r").readlines())

val = parse_mjsynth(path, val)

else:

train = [os.path.join(dp, f) for dp, dn, filenames in os.walk(path)

for f in filenames if re.search('png|jpeg|jpg', f)]

prng.shuffle(train)

length = len(train)

train, val = train[:int(length*train_portion)], train[int(length*train_portion):]

lengths = get_lengths(train)

max_len = max(lengths.values())

print(f' [INFO] {len(train)} train and {len(val)} validation images loaded ')

reader = Readf(

img_size=(imgh, imgW, 1), normed=norm, batch_size=batch_size,

classes=classes, max_len=max_len, transform_p=0.7

)

print(" [INFO] Number of classes: {}; Max. string length: {} ".format(len(classes)+1, max_len))

init_model = CRNN(num_classes=len(classes)+1, shape=(imgh, imgW, 1), GRU=GRU,

time_dense_size=time_dense_size, n_units=n_units, max_string_len=max_len)

model = init_model.get_model()

save_model_json(model, save_path, model_name)

if pretrained_path is not None:

model.load_weights(pretrained_path)

train_steps = len(train) // batch_size

if (len(train) % batch_size) > 0:

train_steps += 1

test_steps = len(val) // batch_size

if (len(val) % batch_size) > 0:

test_steps += 1

start_time = time.time()

with open(save_path+'/'+model_name + '/model_summary.txt','w') as f:

model.summary(print_fn=lambda x: f.write(x + '\n'))

model.summary()

if opt == "adam":

optimizer = optimizers.Adam(lr=lr, beta_1=0.5, beta_2=0.999, clipnorm=5)

elif opt == "sgd":

optimizer = optimizers.SGD(lr=lr, decay=1e-6, momentum=0.9, nesterov=True, clipnorm=5)

model.compile(loss={"ctc": lambda y_true, y_pred: y_pred}, optimizer=optimizer)

callbacks_list = []

callbacks_list.append(ModelCheckpoint(filepath=save_path+'/%s/checkpoint_weights.h5'%model_name, verbose=1,

save_best_only=True, save_weights_only=True))

if early_stopping:

callbacks_list.append(EarlyStoppingIter(monitor='loss', min_delta=.0001, patience=early_stopping,

verbose=1, restore_best_weights=True, mode="auto"))

H = model.fit_generator(

generator=reader.run_generator(train, downsample_factor=2**init_model.pooling_counter_h),

steps_per_epoch=train_steps,

epochs=nbepochs,

validation_data=reader.run_generator(val, downsample_factor=2**init_model.pooling_counter_h),

validation_steps=test_steps,

shuffle=False, verbose=1,

callbacks=callbacks_list

)

pickle.dump(H.history, open(save_path+'/'+model_name+'/loss_history.pickle.dat', 'wb'))

print(" [INFO] Training finished in %i sec.!" % (round(time.time() - start_time, 2)))

model.save_weights(save_path+'/'+model_name+"/final_weights.h5")

model.save(save_path+'/'+model_name+"/final_model.h5")

print(" [INFO] Models and history saved! ")