NoteOnMe博客平台搭建
Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

74 Zeilen
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)