utils.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. import colorsys
  2. from enum import Enum
  3. from functools import reduce
  4. from PIL import ImageDraw, ImageFont, Image
  5. from keras import backend as K
  6. import numpy as np
  7. import cv2
  8. def compose(*funcs):
  9. """Compose arbitrarily many functions, evaluated left to right.
  10. Reference: https://mathieularose.com/function-composition-in-python/
  11. """
  12. if funcs:
  13. return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
  14. else:
  15. raise ValueError('Composition of empty sequence not supported.')
  16. def box_iou(b1, b2):
  17. """Return iou tensor
  18. Parameters
  19. ----------
  20. b1: tensor, shape=(i1,...,iN, 4), xywh
  21. b2: tensor, shape=(j, 4), xywh
  22. Returns
  23. -------
  24. iou: tensor, shape=(i1,...,iN, j)
  25. """
  26. # Expand dim to apply broadcasting.
  27. b1 = K.expand_dims(b1, -2)
  28. b1_xy = b1[..., :2]
  29. b1_wh = b1[..., 2:4]
  30. b1_wh_half = b1_wh/2.
  31. b1_mins = b1_xy - b1_wh_half
  32. b1_maxes = b1_xy + b1_wh_half
  33. # Expand dim to apply broadcasting.
  34. b2 = K.expand_dims(b2, 0)
  35. b2_xy = b2[..., :2]
  36. b2_wh = b2[..., 2:4]
  37. b2_wh_half = b2_wh/2.
  38. b2_mins = b2_xy - b2_wh_half
  39. b2_maxes = b2_xy + b2_wh_half
  40. intersect_mins = K.maximum(b1_mins, b2_mins)
  41. intersect_maxes = K.minimum(b1_maxes, b2_maxes)
  42. intersect_wh = K.maximum(intersect_maxes - intersect_mins, 0.)
  43. intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
  44. b1_area = b1_wh[..., 0] * b1_wh[..., 1]
  45. b2_area = b2_wh[..., 0] * b2_wh[..., 1]
  46. iou = intersect_area / (b1_area + b2_area - intersect_area)
  47. return iou
  48. def get_classes(classes_path):
  49. """loads the classes"""
  50. with open(classes_path) as f:
  51. class_names = f.readlines()
  52. class_names = [c.strip() for c in class_names]
  53. return class_names
  54. def get_anchors(anchors_path):
  55. """loads the anchors from a file"""
  56. with open(anchors_path) as f:
  57. anchors = f.readline()
  58. anchors = [float(x) for x in anchors.split(',')]
  59. return np.array(anchors).reshape(-1, 2)
  60. def get_colors(number, bright=True):
  61. """
  62. Generate random colors for drawing bounding boxes.
  63. To get visually distinct colors, generate them in HSV space then
  64. convert to RGB.
  65. """
  66. if number <= 0:
  67. return []
  68. brightness = 1.0 if bright else 0.7
  69. hsv_tuples = [(x / number, 1., brightness)
  70. for x in range(number)]
  71. colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
  72. colors = list(
  73. map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
  74. colors))
  75. np.random.seed(10101) # Fixed seed for consistent colors across runs.
  76. np.random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
  77. np.random.seed(None) # Reset seed to default.
  78. return colors
  79. labelType = Enum('labelType', ('LABEL_TOP_OUTSIDE',
  80. 'LABEL_BOTTOM_OUTSIDE',
  81. 'LABEL_TOP_INSIDE',
  82. 'LABEL_BOTTOM_INSIDE',))
  83. def draw_boxes(image, out_boxes, out_classes, out_scores, class_names, colors):
  84. print()
  85. font = ImageFont.truetype(font='yolo_data/FiraMono-Medium.otf',
  86. size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
  87. thickness = (image.size[0] + image.size[1]) // 300
  88. box_list = []
  89. for i, c in reversed(list(enumerate(out_classes))):
  90. predicted_class = class_names[c]
  91. box = out_boxes[i]
  92. score = out_scores[i]
  93. label = '{} {:.2f}'.format(predicted_class, score)
  94. draw = ImageDraw.Draw(image)
  95. label_size = draw.textsize(label, font)
  96. top, left, bottom, right = box
  97. top = max(0, np.floor(top + 0.5).astype('int32'))
  98. left = max(0, np.floor(left + 0.5).astype('int32'))
  99. bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
  100. right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
  101. print(label, (left, top), (right, bottom))
  102. box_list.append([(left, top), (right, bottom)])
  103. if top - label_size[1] >= 0:
  104. text_origin = np.array([left, top - label_size[1]])
  105. else:
  106. text_origin = np.array([left, top + 1])
  107. # My kingdom for a good redistributable image drawing library.
  108. for i in range(thickness):
  109. draw.rectangle(
  110. [left + i, top + i, right - i, bottom - i],
  111. outline=colors[c])
  112. draw.rectangle(
  113. [tuple(text_origin), tuple(text_origin + label_size)],
  114. fill=colors[c])
  115. draw.text(text_origin, label, fill=(0, 0, 0), font=font)
  116. del draw
  117. return image, box_list
  118. def draw_label(image, text, color, coords, label_type=labelType.LABEL_TOP_OUTSIDE):
  119. font = cv2.FONT_HERSHEY_PLAIN
  120. font_scale = 1.
  121. (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]
  122. padding = 5
  123. rect_height = text_height + padding * 2
  124. rect_width = text_width + padding * 2
  125. (x, y) = coords
  126. if label_type == labelType.LABEL_TOP_OUTSIDE or label_type == labelType.LABEL_BOTTOM_INSIDE:
  127. cv2.rectangle(image, (x, y), (x + rect_width, y - rect_height), color, cv2.FILLED)
  128. cv2.putText(image, text, (x + padding, y - text_height + padding), font,
  129. fontScale=font_scale,
  130. color=(255, 255, 255),
  131. lineType=cv2.LINE_AA)
  132. else:
  133. # LABEL_BOTTOM_OUTSIDE or LABEL_TOP_INSIDE
  134. cv2.rectangle(image, (x, y), (x + rect_width, y + rect_height), color, cv2.FILLED)
  135. cv2.putText(image, text, (x + padding, y + text_height + padding), font,
  136. fontScale=font_scale,
  137. color=(255, 255, 255),
  138. lineType=cv2.LINE_AA)
  139. return image
  140. def pil_resize(image_np, height, width):
  141. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
  142. image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
  143. image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)
  144. return image_np
  145. def pil_resize_a(image_np, height, width):
  146. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
  147. image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
  148. image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGBA2BGRA)
  149. return image_np
  150. def np2pil(image_np):
  151. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
  152. return image_pil
  153. def pil2np(image_pil):
  154. image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
  155. return image_np
  156. def np2pil_a(image_np):
  157. image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
  158. return image_pil
  159. def pil2np_a(image_pil):
  160. image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGBA2BGRA)
  161. return image_np
  162. def pil_rotate(image_np, angle, fill_color=(255, 255, 255)):
  163. image_pil = np2pil(image_np)
  164. image_rotate = image_pil.rotate(angle, expand=False, fillcolor=fill_color)
  165. image_np = pil2np(image_rotate)
  166. return image_np
  167. def pil_rotate_a(image_np, angle):
  168. image_pil = np2pil_a(image_np)
  169. image_rotate = image_pil.rotate(angle, expand=False, fillcolor=(255, 255, 255))
  170. image_np = pil2np_a(image_rotate)
  171. return image_np