NoteOnMe博客平台搭建
Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.

74 рядки
2.3 KiB

  1. from scipy.misc import imread,imshow
  2. import os
  3. from PIL import Image
  4. import PIL
  5. from model.img2seq import Img2SeqModel
  6. from model.utils.general import Config, run
  7. from model.utils.text import Vocab
  8. from model.utils.image import greyscale, crop_image, pad_image,predictsize ,\
  9. downsample_image, TIMEOUT
  10. def interactive_shell(model):
  11. """Creates interactive shell to play with model
  12. """
  13. model.logger.info("""
  14. This is an interactive mode.
  15. To exit, enter 'exit'.
  16. Enter a path to a file
  17. input> data/images_test/0.png""")
  18. while True:
  19. img_path = input("input> ")
  20. if img_path == "exit":
  21. break
  22. if img_path[-3:] == "png":
  23. img = imread(img_path)
  24. elif img_path[-3:] == "pdf":
  25. # call magick to convert the pdf into a png file
  26. buckets = [
  27. [240, 100], [320, 80], [400, 80], [400, 100], [480, 80], [480, 100],
  28. [560, 80], [560, 100], [640, 80], [640, 100], [720, 80], [720, 100],
  29. [720, 120], [720, 200], [800, 100], [800, 320], [1000, 200],
  30. [1000, 400], [1200, 200], [1600, 200], [1600, 1600]
  31. ]
  32. dir_output = "tmp/"
  33. name = img_path.split('/')[-1].split('.')[0]
  34. run("magick convert -density {} -quality {} {} {}".format(200, 100,
  35. img_path, dir_output+"{}.png".format(name)), TIMEOUT)
  36. img_path = dir_output + "{}.png".format(name)
  37. crop_image(img_path, img_path)
  38. pad_image(img_path, img_path, buckets=buckets)
  39. downsample_image(img_path, img_path, 2)
  40. img = imread(img_path)
  41. img = predictsize(img)
  42. im_converted = PIL.Image.fromarray(img)
  43. im_converted.show()
  44. img = greyscale(img)
  45. hyps = model.predict(img)
  46. with open("norm_formula_val.txt", "w") as f:
  47. f.write(hyps[0])
  48. model.logger.info(hyps[0])
  49. if __name__ == "__main__":
  50. # restore config and model
  51. dir_output = "results/full/"
  52. config_vocab = Config(dir_output + "vocab.json")
  53. config_model = Config(dir_output + "model.json")
  54. vocab = Vocab(config_vocab)
  55. model = Img2SeqModel(config_model, dir_output, vocab)
  56. model.build_pred()
  57. model.restore_session(dir_output + "model.weights4/test-model.ckpt")
  58. interactive_shell(model)