def main(data, vocab, training, model, output):
# Load configs
dir_output = output
config = Config([data, vocab, training, model])
config.save(dir_output)
vocab = Vocab(config)
# Load datasets
train_set = DataGenerator(path_formulas=config.path_formulas_train,
dir_images=config.dir_images_train, img_prepro=greyscale,
max_iter=config.max_iter, bucket=config.bucket_train,
path_matching=config.path_matching_train,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
val_set = DataGenerator(path_formulas=config.path_formulas_val,
dir_images=config.dir_images_val, img_prepro=greyscale,
max_iter=config.max_iter, bucket=config.bucket_val,
path_matching=config.path_matching_val,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
# Define learning rate schedule
n_batches_epoch = ((len(train_set) + config.batch_size - 1) //
config.batch_size)
lr_schedule = LRSchedule(lr_init=config.lr_init,
start_decay=config.start_decay*n_batches_epoch,
end_decay=config.end_decay*n_batches_epoch,
end_warm=config.end_warm*n_batches_epoch,
lr_warm=config.lr_warm,
lr_min=config.lr_min)
# Build model and train
model = Img2SeqModel(config, dir_output, vocab)
model.build_train(config)
model.train(config, train_set, val_set, lr_schedule)
def main(results):
# restore config and model
dir_output = results
config_data = Config(dir_output + "data.json")
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(dir_output + "model.weights/")
# load dataset
test_set = DataGenerator(path_formulas=config_data.path_formulas_test,
dir_images=config_data.dir_images_test, img_prepro=greyscale,
max_iter=config_data.max_iter, bucket=config_data.bucket_test,
path_matching=config_data.path_matching_test,
max_len=config_data.max_length_formula,
form_prepro=vocab.form_prepro,)
# build images from formulas
formula_ref = path.join(dir_output, "formulas_test/ref.txt")
formula_hyp = path.join(dir_output, "formulas_test/hyp_0.txt")
images_ref = path.join(dir_output, "images_test/ref/")
images_test = path.join(dir_output, "images_test/hyp_0/")
build_images(load_formulas(formula_ref), images_ref)
build_images(load_formulas(formula_hyp), images_test)
# score the repositories
scores = score_dirs(images_ref, images_test, greyscale)
msg = " - ".join(["{} {:04.2f}".format(k, v) for k, v in scores.items()])
model.logger.info("- Eval Img: {}".format(msg))
def main(data, vocab):
data_config = Config(data)
# datasets
train_set = DataGenerator(path_formulas=data_config.path_formulas_train,
dir_images=data_config.dir_images_train,
path_matching=data_config.path_matching_train)
test_set = DataGenerator(path_formulas=data_config.path_formulas_test,
dir_images=data_config.dir_images_test,
path_matching=data_config.path_matching_test)
val_set = DataGenerator(path_formulas=data_config.path_formulas_val,
dir_images=data_config.dir_images_val,
path_matching=data_config.path_matching_val)
# produce images and matching files
# train_set.build(buckets=None, n_threads=1)
train_set.build(buckets=None)
test_set.build(buckets=None)
val_set.build(buckets=None)
# train_set.build(buckets=data_config.buckets)
# test_set.build(buckets=data_config.buckets)
# val_set.build(buckets=data_config.buckets)
# p = Augmentor.Pipeline(data_config.dir_images_train)
# p.zoom(probability=1, min_factor=1.1, max_factor=1.5)
# # p.process()
# augmented_images, labels = p.sample(3)
# print labels
# vocab
vocab_config = Config(vocab)
vocab = build_vocab([train_set], min_count=vocab_config.min_count_tok)
write_vocab(vocab, vocab_config.path_vocab)
def main(data, vocab):
data_config = Config(data)
# datasets
train_set = DataGenerator(
path_formulas=data_config.path_formulas_train,
dir_images=data_config.dir_images_train,
path_matching=data_config.path_matching_train)
test_set = DataGenerator(
path_formulas=data_config.path_formulas_test,
dir_images=data_config.dir_images_test,
path_matching=data_config.path_matching_test)
val_set = DataGenerator(
path_formulas=data_config.path_formulas_val,
dir_images=data_config.dir_images_val,
path_matching=data_config.path_matching_val)
# produce images and matching files
train_set.build(buckets=data_config.buckets)
test_set.build(buckets=data_config.buckets)
val_set.build(buckets=data_config.buckets)
# vocab
vocab_config = Config(vocab)
vocab = build_vocab([train_set], min_count=vocab_config.min_count_tok)
write_vocab(vocab, vocab_config.path_vocab)
def main(results):
# restore config and model
dir_output = results
config_data = Config(dir_output + "data.json")
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(dir_output + "model.weights/")
# load dataset
test_set = DataGenerator(
index_file=config_data.index_test,
path_formulas=config_data.path_formulas_test,
dir_images=config_data.dir_images_test,
max_iter=config_data.max_iter,
path_matching=config_data.path_matching_test,
form_prepro=vocab.form_prepro)
# use model to write predictions in files
config_eval = Config({"dir_answers": dir_output + "formulas_test/",
"batch_size": 20})
files, perplexity = model.write_prediction(config_eval, test_set)
formula_ref, formula_hyp = files[0], files[1]
# score the ref and prediction files
scores = score_files(formula_ref, formula_hyp)
scores["perplexity"] = perplexity
msg = " - ".join(["{} {:04.2f}".format(k, v) for k, v in scores.items()])
model.logger.info("- Test Txt: {}".format(msg))
def main(results):
tf.enable_resource_variables()
# restore config and model
dir_output = results
weights_dir = os.path.join(dir_output, 'model.weights/')
t = datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S')
# saved_path = 'saved_' + t
saved_path = 'saved_word'
saved_path = os.path.join(dir_output, saved_path)
config_data = Config(dir_output + "data.json")
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
if not os.path.isdir(saved_path):
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(weights_dir)
model.save_savedmodel(saved_path)
# chkp.print_tensors_in_checkpoint_file(weights_dir, tensor_name='', all_tensors=True)
SAMPLE_DIR = 'tools/data/hand/raw_word'
def representative_dataset_gen():
num_calibration_steps = 10
if not os.path.isdir(SAMPLE_DIR):
print 'Failed to read representative_dataset'
return
for f in os.listdir(SAMPLE_DIR):
img_path = os.path.join(SAMPLE_DIR, f)
img = Image.open(img_path)
img = img.resize((80, 100), Image.BILINEAR)
img.show()
img = np.array(img)
yield [img]
num_calibration_steps -= 1
if num_calibration_steps == 0:
break
converter = tf.lite.TFLiteConverter.from_saved_model(saved_path)
converter.target_ops = [
# tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS
]
# Following has "Segmentation fault"
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.representative_dataset = representative_dataset_gen
tflite_model = converter.convert()
open("converted_model_word.tflite", "wb").write(tflite_model)
def getModelForPrediction():
# restore config and model
dir_output = "./results/full/"
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
# model.restore_session(dir_output + "model_weights/model.cpkt")
return model
def main(data, vocab, training, model, output):
# Load configs
dir_output = output
config = Config([data, vocab, training, model])
config.save(dir_output)
vocab = Vocab(config)
# Load datasets
train_set = DataGenerator(path_formulas=config.path_formulas_train,
dir_images=config.dir_images_train, img_prepro=greyscale,
max_iter=config.max_iter, bucket=config.bucket_train,
path_matching=config.path_matching_train,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
all_img = []
all_formula = []
for i, (_img, _formula) in enumerate(minibatches(train_set, batch_size)):
all_img.append(_img)
if _formula is not None:
_formula, _formula_length = pad_batch_formulas(
_formula,
vocab.id_pad,
vocab.id_end
)
all_formula.append(_formula)
np.save('np_formula', np.array(all_formula))
np.save('np_img', np.array(all_img))
print("DONE EXPORTING NUMPY FILES")
return None
val_set = DataGenerator(path_formulas=config.path_formulas_val,
dir_images=config.dir_images_val, img_prepro=greyscale,
max_iter=config.max_iter, bucket=config.bucket_val,
path_matching=config.path_matching_val,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
# Define learning rate schedule
n_batches_epoch = ((len(train_set) + config.batch_size - 1) //
config.batch_size)
lr_schedule = LRSchedule(lr_init=config.lr_init,
start_decay=config.start_decay*n_batches_epoch,
end_decay=config.end_decay*n_batches_epoch,
end_warm=config.end_warm*n_batches_epoch,
lr_warm=config.lr_warm,
lr_min=config.lr_min)
# Build model and train
model = Img2SeqModel(config, dir_output, vocab)
model.build_train(config)
model.train(config, train_set, val_set, lr_schedule)
def load_model(dir_output="results/full/",
vocab_config='vocab.json',
model_config='model.json',
model_path='model.weights/'):
config_vocab = Config(os.path.join(dir_output, vocab_config))
config_model = Config(os.path.join(dir_output, model_config))
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(os.path.join(dir_output, model_path))
return model
def main(data, vocab, training, model, output):
# Load configs
dir_output = output
config = Config([data, vocab, training, model])
config.save(dir_output)
vocab = Vocab(config)
# Load datasets
train_set = DataGenerator(
index_file=config.index_train,
path_formulas=config.path_formulas_train,
dir_images=config.dir_images_train,
max_iter=config.max_iter,
path_matching=config.path_matching_train,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
val_set = DataGenerator(
index_file=config.index_val,
path_formulas=config.path_formulas_val,
dir_images=config.dir_images_val,
max_iter=config.max_iter,
path_matching=config.path_matching_val,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro)
# Define learning rate schedule
n_batches_epoch = ((len(train_set) + config.batch_size - 1) //
config.batch_size)
print len(train_set)
print config.batch_size
print n_batches_epoch
lr_schedule = LRSchedule(lr_init=config.lr_init,
start_decay=config.start_decay * n_batches_epoch,
end_decay=config.end_decay * n_batches_epoch,
end_warm=config.end_warm * n_batches_epoch,
lr_warm=config.lr_warm,
lr_min=config.lr_min)
transfer_model = config.transfer_model
# Build model and train
model = Img2SeqModel(config, dir_output, vocab)
model.build_train(config)
if transfer_model and os.path.isdir(transfer_model):
model.restore_session(transfer_model)
model.train(config, train_set, val_set, lr_schedule)
def _run_epoch(self, config, train_set, val_set, epoch, lr_schedule):
"""Performs an epoch of training
Args:
config: Config instance
train_set: Dataset instance
val_set: Dataset instance
epoch: (int) id of the epoch, starting at 0
lr_schedule: LRSchedule instance that takes care of learning proc
Returns:
score: (float) model will select weights that achieve the highest
score
"""
# logging
batch_size = config.batch_size
nbatches = (len(train_set) + batch_size - 1) // batch_size
prog = Progbar(nbatches)
# iterate over dataset
for i, (img, formula) in enumerate(minibatches(train_set, batch_size)):
# get feed dict
fd = self._get_feed_dict(img,
training=True,
formula=formula,
lr=lr_schedule.lr,
dropout=config.dropout)
# update step
_, loss_eval = self.sess.run([self.train_op, self.loss],
feed_dict=fd)
prog.update(i + 1, [("loss", loss_eval),
("perplexity", np.exp(loss_eval)),
("lr", lr_schedule.lr)])
# update learning rate
lr_schedule.update(batch_no=epoch * nbatches + i)
# 生成summary
summary_str = self.sess.run(self.merged, feed_dict=fd)
self.file_writer.add_summary(summary_str, epoch) # 将summary 写入文件
# if (i+1) % 100 == 0:
# # 太慢了，读了 100 批次后就保存先，保存的权重要用于调试 attention
# self.save_debug_session(epoch, i)
# logging
self.logger.info("- Training: {}".format(prog.info))
# evaluation
config_eval = Config({
"dir_answers": self._dir_output + "formulas_val/",
"batch_size": config.batch_size
})
scores = self.evaluate(config_eval, val_set)
score = scores[config.metric_val]
lr_schedule.update(score=score)
return score
def get_im2latex_model(weight_dir):
"""
Load up model from the given weight location
:param weight_dir: weight location
:return: trained model
"""
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
tf.logging.set_verbosity(tf.logging.ERROR)
config_vocab = Config(weight_dir + "vocab.json")
config_model = Config(weight_dir + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, weight_dir, vocab)
model.build_pred()
model.restore_session(weight_dir + "model.weights/")
return model
def _run_train_epoch(self, config, train_set, val_set, epoch, lr_schedule):
"""Performs an epoch of training
Args:
config: Config instance
train_set: Dataset instance
val_set: Dataset instance
epoch: (int) id of the epoch, starting at 0
lr_schedule: LRSchedule instance that takes care of learning proc
Returns:
score: (float) model will select weights that achieve the highest score
"""
# logging
batch_size = config.batch_size
nbatches = (len(train_set) + batch_size - 1) // batch_size
prog = Progbar(nbatches)
self.model.train()
self.encoder.train()
self.decoder.train()
train_loader = torch.utils.data.DataLoader(
ImgFormulaDataset(train_set),
batch_size=batch_size,
shuffle=True,
num_workers=3,
pin_memory=True)
# for i, (img, formula) in enumerate(train_loader):
for i, (img, formula) in enumerate(minibatches(train_set, batch_size)):
img = pad_batch_images_2(img)
img = torch.FloatTensor(img) # (N, W, H, C)
formula, formula_length = pad_batch_formulas(
formula, self._vocab.id_pad, self._vocab.id_end)
img = img.permute(0, 3, 1, 2) # (N, C, W, H)
formula = torch.LongTensor(formula) # (N,)
loss_eval = self.getLoss(img,
formula=formula,
lr=lr_schedule.lr,
dropout=config.dropout,
training=True)
prog.update(i + 1, [("loss", loss_eval), ("lr", lr_schedule.lr)])
# update learning rate
lr_schedule.update(batch_no=epoch * nbatches + i)
self.logger.info("- Training: {}".format(prog.info))
self.logger.info("- Config: (before evaluate, we need to see config)")
config.show(fun=self.logger.info)
# evaluation
config_eval = Config({
"dir_answers": self._dir_output + "formulas_val/",
"batch_size": config.batch_size
})
scores = self.evaluate(config_eval, val_set)
score = scores["perplexity"]
lr_schedule.update(score=score)
return score
def main(image, vocab, model, output):
dir_output = output
img_path = image
# 加载配置，根据配置初始化字典和模型
config = Config([vocab, model])
vocab = Vocab(config)
img2SeqModel = Img2SeqModel(config, dir_output, vocab)
img2SeqModel.build_pred()
vis_img_with_attention(img2SeqModel, img_path, dir_output)
def img2latex_api(weight_dir, img_path, downsample_image_ratio, cropping,
padding, gray_scale):
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.logging.set_verbosity(tf.logging.ERROR)
config_vocab = Config(weight_dir + "vocab.json")
config_model = Config(weight_dir + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, weight_dir, vocab)
model.build_pred()
model.restore_session(weight_dir + "model.weights/")
seq = iaa.Sequential([iaa.GammaContrast(2)])
latex, _, _ = img2latex(model,
img_path,
downsample_image_ratio=downsample_image_ratio,
cropping=cropping,
padding=padding,
img_augment=seq,
gray_scale=gray_scale)
processed_latex = postprocess(latex)
return processed_latex
def main(results):
# restore config and model
dir_output = results
weights_dir = os.path.join(dir_output, 'model.weights/')
saved_path = os.path.join(dir_output, 'saved')
config_data = Config(dir_output + "data.json")
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
if not os.path.isdir(saved_path):
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(weights_dir)
model.save_savedmodel(saved_path)
# chkp.print_tensors_in_checkpoint_file(weights_dir, tensor_name='', all_tensors=True)
converter = tf.lite.TFLiteConverter.from_saved_model(saved_path)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)
def predict(trained_model):
"""
Args:
trained_model (Img2SeqModel):
"""
if len(sys.argv) <= 1:
print "Usage"
print "python predict.py abs/path/to/image.png"
return
image_path = sys.argv[1]
img = imread(image_path)
img = greyscale(img)
res = trained_model.predict(img)
trained_model.logger.info(res[0])
if __name__ == "__main__":
# restore config and model
dir_output = "results/full/"
config_vocab = Config(dir_output + "vocab.json")
config_model = Config(dir_output + "model.json")
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, dir_output, vocab)
model.build_pred()
model.restore_session(dir_output + "model.weights/")
# interactive_shell(model)
predict(model)
from model.utils.general import Config, Progbar, minibatches
# from model.evaluation.text import score_files, write_answers, truncate_end
# _attn_cell_config = {
# 'cell_type': 'lstm',
# 'num_units': 12,
# 'dim_e' : 14,
# 'dim_o' : 16,
# 'dim_embeddings': 32,
# }
dir_output = "results/small/"
config = Config([
"configs/data_small.json",
"configs/vocab_small.json",
"configs/training_small.json",
"configs/model.json",
])
config.save(dir_output)
vocab = Vocab(config)
train_set = DataGenerator(
path_formulas=config.path_formulas_train,
dir_images=config.dir_images_train,
img_prepro=greyscale,
max_iter=config.max_iter,
bucket=config.bucket_train,
path_matching=config.path_matching_train,
max_len=config.max_length_formula,
form_prepro=vocab.form_prepro
)
def _run_train(self, config, train_set, val_set, epoch, lr_schedule):
"""Performs an epoch of training
Args:
config: Config instance
train_set: Dataset instance
val_set: Dataset instance
epoch: (int) id of the epoch, starting at 0
lr_schedule: LRSchedule instance that takes care of learning proc
Returns:
score: (float) model will select weights that achieve the highest score
"""
# logging
batch_size = config.batch_size
train_set.shuffle()
nbatches = (len(train_set) + batch_size - 1) // batch_size
prog = Progbar(nbatches)
# iterate over dataset
for i, (img, formula) in enumerate(minibatches(train_set, batch_size)):
# get feed dict
fd = self._get_feed_dict(img, formula=formula, lr=lr_schedule.lr, dropout=config.dropout)
# 来试试随机的 dropout
# random_dropout = 0.5 + random.random() * 0.5
# fd = self._get_feed_dict(img, formula=formula, lr=lr_schedule.lr, dropout=random_dropout)
# update step
_, G_loss_eval = self.sess.run([self.train_op, self.loss], feed_dict=fd)
if (G_loss_eval <= 1):
# 等 Generator 的 loss 下降到 1 以下再训练判别器，不然判别器训练了感觉都没啥意义
_, D_loss_eval = self.sess.run([self.D_optimizer, self.D_loss], feed_dict=fd)
else:
D_loss_eval = 0
prog.update(i + 1, [("D_loss", D_loss_eval),
("G_loss", G_loss_eval),
("G_perplexity", np.exp(G_loss_eval)),
("lr", lr_schedule.lr)])
# update learning rate
lr_schedule.update(batch_no=epoch*nbatches + i)
# 生成summary
if (i+1) % 10 == 0:
summary_str = self.sess.run(self.merged, feed_dict=fd)
self.file_writer.add_summary(summary_str, epoch) # 将summary 写入文件
# if (i+1) % 100 == 0:
# # 太慢了，读了 100 批次后就保存先，保存的权重要用于调试 attention
# self.save_debug_session(epoch, i)
# logging
self.logger.info("- Training: {}".format(prog.info))
# evaluation
config_eval = Config({
"dir_answers": self._dir_output + "formulas_val/",
"batch_size": config.batch_size
})
scores = self.evaluate(config_eval, val_set)
score = scores["perplexity"] + (scores["ExactMatchScore"] + scores["BLEU-4"] + scores["EditDistance"]) / 10
lr_schedule.update(score=score)
return score
import os
from model.img2seq import Img2SeqModel
from model.utils.general import Config, run
from model.utils.text import Vocab
APP_ROOT = os.path.dirname(
os.path.abspath(__file__)) # refers to application_top
MODEL_FOLDER = os.path.join(APP_ROOT, 'full')
print("loading model ...")
config_vocab = Config(os.path.join(MODEL_FOLDER, "vocab.json"))
config_model = Config(os.path.join(MODEL_FOLDER, "model.json"))
vocab = Vocab(config_vocab)
model = Img2SeqModel(config_model, "/tmp/mxlatexfull", vocab)
model.build_pred()
model.restore_session(os.path.join(MODEL_FOLDER, "model.weights/"))
