You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
2.3 KiB
74 lines
2.3 KiB
from scipy.misc import imread,imshow
|
|
import os
|
|
from PIL import Image
|
|
import PIL
|
|
|
|
from model.img2seq import Img2SeqModel
|
|
from model.utils.general import Config, run
|
|
from model.utils.text import Vocab
|
|
from model.utils.image import greyscale, crop_image, pad_image,predictsize ,\
|
|
downsample_image, TIMEOUT
|
|
|
|
|
|
|
|
def interactive_shell(model):
|
|
"""Creates interactive shell to play with model
|
|
"""
|
|
model.logger.info("""
|
|
This is an interactive mode.
|
|
To exit, enter 'exit'.
|
|
Enter a path to a file
|
|
input> data/images_test/0.png""")
|
|
|
|
while True:
|
|
|
|
img_path = input("input> ")
|
|
|
|
if img_path == "exit":
|
|
break
|
|
|
|
if img_path[-3:] == "png":
|
|
img = imread(img_path)
|
|
|
|
elif img_path[-3:] == "pdf":
|
|
# call magick to convert the pdf into a png file
|
|
buckets = [
|
|
[240, 100], [320, 80], [400, 80], [400, 100], [480, 80], [480, 100],
|
|
[560, 80], [560, 100], [640, 80], [640, 100], [720, 80], [720, 100],
|
|
[720, 120], [720, 200], [800, 100], [800, 320], [1000, 200],
|
|
[1000, 400], [1200, 200], [1600, 200], [1600, 1600]
|
|
]
|
|
|
|
dir_output = "tmp/"
|
|
name = img_path.split('/')[-1].split('.')[0]
|
|
run("magick convert -density {} -quality {} {} {}".format(200, 100,
|
|
img_path, dir_output+"{}.png".format(name)), TIMEOUT)
|
|
img_path = dir_output + "{}.png".format(name)
|
|
crop_image(img_path, img_path)
|
|
pad_image(img_path, img_path, buckets=buckets)
|
|
downsample_image(img_path, img_path, 2)
|
|
|
|
img = imread(img_path)
|
|
|
|
img = predictsize(img)
|
|
im_converted = PIL.Image.fromarray(img)
|
|
im_converted.show()
|
|
img = greyscale(img)
|
|
|
|
hyps = model.predict(img)
|
|
with open("norm_formula_val.txt", "w") as f:
|
|
f.write(hyps[0])
|
|
model.logger.info(hyps[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.weights4/test-model.ckpt")
|
|
interactive_shell(model)
|