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.

utils.py 6.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. """Utilities functions."""
  2. import json
  3. import logging
  4. import os
  5. import sys
  6. import zipfile
  7. from re import finditer
  8. import colorama
  9. import coloredlogs
  10. import cv2
  11. import imageio
  12. import numpy as np
  13. import requests
  14. from PIL import Image
  15. from config import Config as Conf
  16. def read_image(path):
  17. """
  18. Read a file image.
  19. :param path: <string> Path of the image
  20. :return: <RGB> image
  21. """
  22. # Read image
  23. with open(path, "rb") as file:
  24. image_bytes = bytearray(file.read())
  25. np_image = np.asarray(image_bytes, dtype=np.uint8)
  26. image = cv2.imdecode(np_image, cv2.IMREAD_COLOR)
  27. # See if image loaded correctly
  28. if image is None:
  29. Conf.log.error("{} file is not valid image".format(path))
  30. sys.exit(1)
  31. return image
  32. def write_image(image, path):
  33. """
  34. Write a file image to the path (create the directory if needed).
  35. :param image: <RGB> image to write
  36. :param path: <string> location where write the image
  37. :return: None
  38. """
  39. dir_path = os.path.dirname(path)
  40. if dir_path != '':
  41. os.makedirs(dir_path, exist_ok=True)
  42. if os.path.splitext(path)[1] not in cv2_supported_extension():
  43. Conf.log.error("{} invalid extension format.".format(path))
  44. sys.exit(1)
  45. cv2.imwrite(path, image)
  46. if not check_image_file_validity(path):
  47. Conf.log.error(
  48. "Something gone wrong writing {} image file. The final result is not a valid image file.".format(path))
  49. sys.exit(1)
  50. def check_shape(path, shape=Conf.desired_shape):
  51. """
  52. Validate the shape of an image.
  53. :param image: <RGB> Image to check
  54. :param shape: <(int,int,int)> Valid shape
  55. :return: None
  56. """
  57. if os.path.splitext(path)[1] != ".gif":
  58. img_shape = read_image(path).shape
  59. else:
  60. img_shape = imageio.mimread(path)[0][:, :, :3].shape
  61. if img_shape != shape:
  62. Conf.log.error("{} Image is not 512 x 512, got shape: {}".format(path, img_shape))
  63. Conf.log.error("You should use one of the rescale options or manually resize the image")
  64. sys.exit(1)
  65. def check_image_file_validity(image_path):
  66. """
  67. Check is a file is valid image file.
  68. :param image_path: <string> Path to the file to check
  69. :return: <Boolean> True if it's an image file
  70. """
  71. try:
  72. im = Image.open(image_path)
  73. im.verify()
  74. except Exception:
  75. return False
  76. return True if os.stat(image_path).st_size != 0 else False
  77. def setup_log(log_lvl=logging.INFO):
  78. """
  79. Configure a logger.
  80. :param log_lvl: <loggin.LVL> level of the log
  81. :return: <Logger> a logger
  82. """
  83. colorama.init()
  84. log = logging.getLogger(__name__)
  85. coloredlogs.install(level=log_lvl, fmt='[%(levelname)s] %(message)s')
  86. return log
  87. def camel_case_to_str(identifier):
  88. """
  89. Return the string representation of a Camel case word.
  90. :param identifier: <string> camel case word
  91. :return: a string representation
  92. """
  93. matches = finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)', identifier)
  94. return " ".join([m.group(0) for m in matches])
  95. def cv2_supported_extension():
  96. """
  97. List of extension supported by cv2.
  98. :return: <string[]> extensions list
  99. """
  100. return [".bmp", ".dib", ".jpeg", ".jpg", ".jpe", ".jp2", ".png",
  101. ".pbm", ".pgm", "ppm", ".sr", ".ras", ".tiff", ".tif",
  102. ".BMP", ".DIB", ".JPEG", ".JPG", ".JPE", ".JP2", ".PNG",
  103. ".PBM", ".PGM", "PPM", ".SR", ".RAS", ".TIFF", ".TIF"]
  104. def load_json(a):
  105. """
  106. Load a json form file or string.
  107. :param a: <string> Path of the file to load or a json string
  108. :return: <dict> json structure
  109. """
  110. if os.path.isfile(a):
  111. with open(a, 'r') as f:
  112. j = json.load(f)
  113. else:
  114. j = json.loads(str(a))
  115. return j
  116. def json_to_argv(data):
  117. """
  118. Json to args parameters.
  119. :param data: <json>
  120. :return: <Dict>
  121. """
  122. argv = []
  123. for k, v in data.items():
  124. if not isinstance(v, bool):
  125. argv.extend(["--{}".format(k), str(v)])
  126. elif v:
  127. argv.append("--{}".format(k))
  128. return argv
  129. def dl_file(url, file_path):
  130. """
  131. Download a file.
  132. :param url: <string> url of the file to download
  133. :param file_path: <string> file path where save the file
  134. :return: <string> full path of downloaded file
  135. """
  136. Conf.log.debug("Download url : {} to path: {}".format(url, file_path))
  137. response = requests.get(url, stream=True)
  138. dir_path = os.path.dirname(file_path)
  139. if dir_path != '':
  140. os.makedirs(dir_path, exist_ok=True)
  141. with open(file_path, "wb") as f:
  142. total_length = response.headers.get('content-length')
  143. if total_length is None: # no content length header
  144. f.write(response.content)
  145. else:
  146. dl = 0
  147. total_length = int(total_length)
  148. for data in response.iter_content(chunk_size=4096):
  149. dl += len(data)
  150. f.write(data)
  151. done = int(50 * dl / total_length)
  152. print("[{}{}]".format('=' * done, ' ' * (50 - done)), end="\r")
  153. print(" " * 80, end="\r")
  154. return file_path
  155. def unzip(zip_path, extract_path):
  156. """
  157. Extract a zip.
  158. :param zip_path: <string> path to zip to extract
  159. :param extract_path: <string> path to dir where to extract
  160. :return: None
  161. """
  162. Conf.log.debug("Extracting zip : {} to path: {}".format(zip_path, extract_path))
  163. if not os.path.exists(extract_path):
  164. os.makedirs(extract_path, exist_ok=True)
  165. with zipfile.ZipFile(zip_path, "r") as zf:
  166. uncompress_size = sum((file.file_size for file in zf.infolist()))
  167. extracted_size = 0
  168. for file in zf.infolist():
  169. done = int(50 * extracted_size / uncompress_size)
  170. print("[{}{}]".format('=' * done, ' ' * (50 - done)), end="\r")
  171. zf.extract(file, path=extract_path)
  172. extracted_size += file.file_size
  173. print(" " * 80, end="\r")
  174. def is_a_supported_image_file_extension(path):
  175. """
  176. Return true if the file is an image file supported extensions.
  177. :param path: <sting> path of the file to check
  178. :return: <boolean> True if the extension is supported
  179. """
  180. return os.path.splitext(path)[1] in cv2_supported_extension() + [".gif"]
  181. def check_url(url):
  182. """
  183. Check if a url exists withtout downloading it
  184. :return: <boolean> True if return url exists
  185. """
  186. resp = requests.head(url)
  187. return resp.status_code < 400