utils.py 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276
  1. import colorsys
  2. import socket
  3. import traceback
  4. from enum import Enum
  5. from functools import reduce
  6. import requests
  7. from PIL import ImageDraw, ImageFont, Image
  8. from keras import backend as K
  9. import numpy as np
  10. import cv2
  11. def compose(*funcs):
  12. """Compose arbitrarily many functions, evaluated left to right.
  13. Reference: https://mathieularose.com/function-composition-in-python/
  14. """
  15. if funcs:
  16. return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
  17. else:
  18. raise ValueError('Composition of empty sequence not supported.')
  19. def box_iou(b1, b2):
  20. """Return iou tensor
  21. Parameters
  22. ----------
  23. b1: tensor, shape=(i1,...,iN, 4), xywh
  24. b2: tensor, shape=(j, 4), xywh
  25. Returns
  26. -------
  27. iou: tensor, shape=(i1,...,iN, j)
  28. """
  29. # Expand dim to apply broadcasting.
  30. b1 = K.expand_dims(b1, -2)
  31. b1_xy = b1[..., :2]
  32. b1_wh = b1[..., 2:4]
  33. b1_wh_half = b1_wh/2.
  34. b1_mins = b1_xy - b1_wh_half
  35. b1_maxes = b1_xy + b1_wh_half
  36. # Expand dim to apply broadcasting.
  37. b2 = K.expand_dims(b2, 0)
  38. b2_xy = b2[..., :2]
  39. b2_wh = b2[..., 2:4]
  40. b2_wh_half = b2_wh/2.
  41. b2_mins = b2_xy - b2_wh_half
  42. b2_maxes = b2_xy + b2_wh_half
  43. intersect_mins = K.maximum(b1_mins, b2_mins)
  44. intersect_maxes = K.minimum(b1_maxes, b2_maxes)
  45. intersect_wh = K.maximum(intersect_maxes - intersect_mins, 0.)
  46. intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
  47. b1_area = b1_wh[..., 0] * b1_wh[..., 1]
  48. b2_area = b2_wh[..., 0] * b2_wh[..., 1]
  49. iou = intersect_area / (b1_area + b2_area - intersect_area)
  50. return iou
  51. def get_classes(classes_path):
  52. """loads the classes"""
  53. with open(classes_path) as f:
  54. class_names = f.readlines()
  55. class_names = [c.strip() for c in class_names]
  56. return class_names
  57. def get_anchors(anchors_path):
  58. """loads the anchors from a file"""
  59. with open(anchors_path) as f:
  60. anchors = f.readline()
  61. anchors = [float(x) for x in anchors.split(',')]
  62. return np.array(anchors).reshape(-1, 2)
  63. def get_colors(number, bright=True):
  64. """
  65. Generate random colors for drawing bounding boxes.
  66. To get visually distinct colors, generate them in HSV space then
  67. convert to RGB.
  68. """
  69. if number <= 0:
  70. return []
  71. brightness = 1.0 if bright else 0.7
  72. hsv_tuples = [(x / number, 1., brightness)
  73. for x in range(number)]
  74. colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
  75. colors = list(
  76. map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
  77. colors))
  78. np.random.seed(10101) # Fixed seed for consistent colors across runs.
  79. np.random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
  80. np.random.seed(None) # Reset seed to default.
  81. return colors
  82. labelType = Enum('labelType', ('LABEL_TOP_OUTSIDE',
  83. 'LABEL_BOTTOM_OUTSIDE',
  84. 'LABEL_TOP_INSIDE',
  85. 'LABEL_BOTTOM_INSIDE',))
  86. def draw_boxes(image, out_boxes, out_classes, out_scores, class_names, colors):
  87. font = ImageFont.truetype(font='yolo_data/FiraMono-Medium.otf',
  88. size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
  89. thickness = (image.size[0] + image.size[1]) // 300
  90. box_list = []
  91. for i, c in reversed(list(enumerate(out_classes))):
  92. predicted_class = class_names[c]
  93. box = out_boxes[i]
  94. score = out_scores[i]
  95. label = '{} {:.2f}'.format(predicted_class, score)
  96. draw = ImageDraw.Draw(image)
  97. label_size = draw.textsize(label, font)
  98. top, left, bottom, right = box
  99. top = max(0, np.floor(top + 0.5).astype('int32'))
  100. left = max(0, np.floor(left + 0.5).astype('int32'))
  101. bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
  102. right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
  103. # print(label, (left, top), (right, bottom))
  104. box_list.append([(left, top), (right, bottom)])
  105. if top - label_size[1] >= 0:
  106. text_origin = np.array([left, top - label_size[1]])
  107. else:
  108. text_origin = np.array([left, top + 1])
  109. # My kingdom for a good redistributable image drawing library.
  110. for i in range(thickness):
  111. draw.rectangle(
  112. [left + i, top + i, right - i, bottom - i],
  113. outline=colors[c])
  114. draw.rectangle(
  115. [tuple(text_origin), tuple(text_origin + label_size)],
  116. fill=colors[c])
  117. draw.text(text_origin, label, fill=(0, 0, 0), font=font)
  118. del draw
  119. return image, box_list
  120. def draw_label(image, text, color, coords, label_type=labelType.LABEL_TOP_OUTSIDE):
  121. font = cv2.FONT_HERSHEY_PLAIN
  122. font_scale = 1.
  123. (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]
  124. padding = 5
  125. rect_height = text_height + padding * 2
  126. rect_width = text_width + padding * 2
  127. (x, y) = coords
  128. if label_type == labelType.LABEL_TOP_OUTSIDE or label_type == labelType.LABEL_BOTTOM_INSIDE:
  129. cv2.rectangle(image, (x, y), (x + rect_width, y - rect_height), color, cv2.FILLED)
  130. cv2.putText(image, text, (x + padding, y - text_height + padding), font,
  131. fontScale=font_scale,
  132. color=(255, 255, 255),
  133. lineType=cv2.LINE_AA)
  134. else:
  135. # LABEL_BOTTOM_OUTSIDE or LABEL_TOP_INSIDE
  136. cv2.rectangle(image, (x, y), (x + rect_width, y + rect_height), color, cv2.FILLED)
  137. cv2.putText(image, text, (x + padding, y + text_height + padding), font,
  138. fontScale=font_scale,
  139. color=(255, 255, 255),
  140. lineType=cv2.LINE_AA)
  141. return image
  142. def pil_resize(image_np, height, width):
  143. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
  144. image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
  145. image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)
  146. return image_np
  147. def pil_resize_a(image_np, height, width):
  148. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
  149. image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
  150. image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGBA2BGRA)
  151. return image_np
  152. def np2pil(image_np):
  153. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
  154. return image_pil
  155. def pil2np(image_pil):
  156. image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
  157. return image_np
  158. def np2pil_a(image_np):
  159. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
  160. return image_pil
  161. def pil2np_a(image_pil):
  162. image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGBA2BGRA)
  163. return image_np
  164. def pil_rotate(image_np, angle, fill_color=(255, 255, 255)):
  165. image_pil = np2pil(image_np)
  166. image_rotate = image_pil.rotate(angle, expand=False, fillcolor=fill_color)
  167. image_np = pil2np(image_rotate)
  168. return image_np
  169. def pil_rotate_a(image_np, angle):
  170. image_pil = np2pil_a(image_np)
  171. image_rotate = image_pil.rotate(angle, expand=False, fillcolor=(255, 255, 255))
  172. image_np = pil2np_a(image_rotate)
  173. return image_np
  174. def request_post(url, param, time_out=1000, use_zlib=False):
  175. fails = 0
  176. text = None
  177. while True:
  178. try:
  179. if fails >= 1:
  180. break
  181. headers = {'content-type': 'application/json'}
  182. # result = requests.post(url, data=param, timeout=time_out)
  183. session = requests.Session()
  184. result = session.post(url, data=param, timeout=time_out)
  185. if result.status_code == 200:
  186. text = result.text
  187. break
  188. else:
  189. print('result.status_code', result.status_code)
  190. print('result.text', result.text)
  191. fails += 1
  192. continue
  193. except socket.timeout:
  194. fails += 1
  195. print('timeout! fail times:', fails)
  196. except:
  197. fails += 1
  198. print('fail! fail times:', fails)
  199. traceback.print_exc()
  200. return text
  201. def np2bytes(image_np):
  202. # numpy转为可序列化的string
  203. success, img_encode = cv2.imencode(".jpg", image_np)
  204. # numpy -> bytes
  205. img_bytes = img_encode.tobytes()
  206. return img_bytes
  207. def bytes2np(_b):
  208. try:
  209. # 二进制数据流转np.ndarray [np.uint8: 8位像素]
  210. image_np = cv2.imdecode(np.frombuffer(_b, np.uint8), cv2.IMREAD_COLOR)
  211. # 将rgb转为bgr
  212. # image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
  213. return image_np
  214. except cv2.error as e:
  215. if "src.empty()" in str(e):
  216. print("bytes2np image is empty!")
  217. return None
  218. except:
  219. traceback.print_exc()
  220. return None