utility.py 3.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import logging
  15. import os
  16. import imghdr
  17. import cv2
  18. def print_dict(d, logger, delimiter=0):
  19. """
  20. Recursively visualize a dict and
  21. indenting acrrording by the relationship of keys.
  22. """
  23. for k, v in sorted(d.items()):
  24. if isinstance(v, dict):
  25. logger.info("{}{} : ".format(delimiter * " ", str(k)))
  26. print_dict(v, logger, delimiter + 4)
  27. elif isinstance(v, list) and len(v) >= 1 and isinstance(v[0], dict):
  28. logger.info("{}{} : ".format(delimiter * " ", str(k)))
  29. for value in v:
  30. print_dict(value, logger, delimiter + 4)
  31. else:
  32. logger.info("{}{} : {}".format(delimiter * " ", k, v))
  33. def get_check_global_params(mode):
  34. check_params = ['use_gpu', 'max_text_length', 'image_shape', \
  35. 'image_shape', 'character_type', 'loss_type']
  36. if mode == "train_eval":
  37. check_params = check_params + [ \
  38. 'train_batch_size_per_card', 'test_batch_size_per_card']
  39. elif mode == "test":
  40. check_params = check_params + ['test_batch_size_per_card']
  41. return check_params
  42. def get_image_file_list(img_file):
  43. imgs_lists = []
  44. if img_file is None or not os.path.exists(img_file):
  45. raise Exception("not found any img file in {}".format(img_file))
  46. img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif', 'GIF'}
  47. if os.path.isfile(img_file) and imghdr.what(img_file) in img_end:
  48. imgs_lists.append(img_file)
  49. elif os.path.isdir(img_file):
  50. for single_file in os.listdir(img_file):
  51. file_path = os.path.join(img_file, single_file)
  52. if os.path.isfile(file_path) and imghdr.what(file_path) in img_end:
  53. imgs_lists.append(file_path)
  54. if len(imgs_lists) == 0:
  55. raise Exception("not found any img file in {}".format(img_file))
  56. return imgs_lists
  57. def check_and_read_gif(img_path):
  58. if os.path.basename(img_path)[-3:] in ['gif', 'GIF']:
  59. gif = cv2.VideoCapture(img_path)
  60. ret, frame = gif.read()
  61. if not ret:
  62. logger = logging.getLogger('ppocr')
  63. logger.info("Cannot read {}. This gif image maybe corrupted.")
  64. return None, False
  65. if len(frame.shape) == 2 or frame.shape[-1] == 1:
  66. frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
  67. imgvalue = frame[:, :, ::-1]
  68. return imgvalue, True
  69. return None, False