visualizer.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292
  1. # Copyright (c) 2019 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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. from __future__ import unicode_literals
  18. import numpy as np
  19. from PIL import Image, ImageDraw
  20. import cv2
  21. import os
  22. import math
  23. from .colormap import colormap
  24. from paddlex.ppdet.utils.logger import setup_logger
  25. logger = setup_logger(__name__)
  26. __all__ = ['visualize_results']
  27. def visualize_results(image,
  28. bbox_res,
  29. mask_res,
  30. segm_res,
  31. keypoint_res,
  32. im_id,
  33. catid2name,
  34. threshold=0.5):
  35. """
  36. Visualize bbox and mask results
  37. """
  38. if bbox_res is not None:
  39. image = draw_bbox(image, im_id, catid2name, bbox_res, threshold)
  40. if mask_res is not None:
  41. image = draw_mask(image, im_id, mask_res, threshold)
  42. if segm_res is not None:
  43. image = draw_segm(image, im_id, catid2name, segm_res, threshold)
  44. if keypoint_res is not None:
  45. image = draw_pose(image, keypoint_res, threshold)
  46. return image
  47. def draw_mask(image, im_id, segms, threshold, alpha=0.7):
  48. """
  49. Draw mask on image
  50. """
  51. mask_color_id = 0
  52. w_ratio = .4
  53. color_list = colormap(rgb=True)
  54. img_array = np.array(image).astype('float32')
  55. for dt in np.array(segms):
  56. if im_id != dt['image_id']:
  57. continue
  58. segm, score = dt['segmentation'], dt['score']
  59. if score < threshold:
  60. continue
  61. import pycocotools.mask as mask_util
  62. mask = mask_util.decode(segm) * 255
  63. color_mask = color_list[mask_color_id % len(color_list), 0:3]
  64. mask_color_id += 1
  65. for c in range(3):
  66. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  67. idx = np.nonzero(mask)
  68. img_array[idx[0], idx[1], :] *= 1.0 - alpha
  69. img_array[idx[0], idx[1], :] += alpha * color_mask
  70. return Image.fromarray(img_array.astype('uint8'))
  71. def draw_bbox(image, im_id, catid2name, bboxes, threshold):
  72. """
  73. Draw bbox on image
  74. """
  75. draw = ImageDraw.Draw(image)
  76. catid2color = {}
  77. color_list = colormap(rgb=True)[:40]
  78. for dt in np.array(bboxes):
  79. if im_id != dt['image_id']:
  80. continue
  81. catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
  82. if score < threshold:
  83. continue
  84. if catid not in catid2color:
  85. idx = np.random.randint(len(color_list))
  86. catid2color[catid] = color_list[idx]
  87. color = tuple(catid2color[catid])
  88. # draw bbox
  89. if len(bbox) == 4:
  90. # draw bbox
  91. xmin, ymin, w, h = bbox
  92. xmax = xmin + w
  93. ymax = ymin + h
  94. draw.line(
  95. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
  96. (xmin, ymin)],
  97. width=2,
  98. fill=color)
  99. elif len(bbox) == 8:
  100. x1, y1, x2, y2, x3, y3, x4, y4 = bbox
  101. draw.line(
  102. [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)],
  103. width=2,
  104. fill=color)
  105. xmin = min(x1, x2, x3, x4)
  106. ymin = min(y1, y2, y3, y4)
  107. else:
  108. logger.error('the shape of bbox must be [M, 4] or [M, 8]!')
  109. # draw label
  110. text = "{} {:.2f}".format(catid2name[catid], score)
  111. tw, th = draw.textsize(text)
  112. draw.rectangle(
  113. [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
  114. draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
  115. return image
  116. def save_result(save_path, results, catid2name, threshold):
  117. """
  118. save result as txt
  119. """
  120. img_id = int(results["im_id"])
  121. with open(save_path, 'w') as f:
  122. if "bbox_res" in results:
  123. for dt in results["bbox_res"]:
  124. catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
  125. if score < threshold:
  126. continue
  127. # each bbox result as a line
  128. # for rbox: classname score x1 y1 x2 y2 x3 y3 x4 y4
  129. # for bbox: classname score x1 y1 w h
  130. bbox_pred = '{} {} '.format(catid2name[catid],
  131. score) + ' '.join(
  132. [str(e) for e in bbox])
  133. f.write(bbox_pred + '\n')
  134. elif "keypoint_res" in results:
  135. for dt in results["keypoint_res"]:
  136. kpts = dt['keypoints']
  137. scores = dt['score']
  138. keypoint_pred = [img_id, scores, kpts]
  139. print(keypoint_pred, file=f)
  140. else:
  141. print("No valid results found, skip txt save")
  142. def draw_segm(image,
  143. im_id,
  144. catid2name,
  145. segms,
  146. threshold,
  147. alpha=0.7,
  148. draw_box=True):
  149. """
  150. Draw segmentation on image
  151. """
  152. mask_color_id = 0
  153. w_ratio = .4
  154. color_list = colormap(rgb=True)
  155. img_array = np.array(image).astype('float32')
  156. for dt in np.array(segms):
  157. if im_id != dt['image_id']:
  158. continue
  159. segm, score, catid = dt['segmentation'], dt['score'], dt['category_id']
  160. if score < threshold:
  161. continue
  162. import pycocotools.mask as mask_util
  163. mask = mask_util.decode(segm) * 255
  164. color_mask = color_list[mask_color_id % len(color_list), 0:3]
  165. mask_color_id += 1
  166. for c in range(3):
  167. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  168. idx = np.nonzero(mask)
  169. img_array[idx[0], idx[1], :] *= 1.0 - alpha
  170. img_array[idx[0], idx[1], :] += alpha * color_mask
  171. if not draw_box:
  172. center_y, center_x = ndimage.measurements.center_of_mass(mask)
  173. label_text = "{}".format(catid2name[catid])
  174. vis_pos = (max(int(center_x) - 10, 0), int(center_y))
  175. cv2.putText(img_array, label_text, vis_pos,
  176. cv2.FONT_HERSHEY_COMPLEX, 0.3, (255, 255, 255))
  177. else:
  178. mask = mask_util.decode(segm) * 255
  179. sum_x = np.sum(mask, axis=0)
  180. x = np.where(sum_x > 0.5)[0]
  181. sum_y = np.sum(mask, axis=1)
  182. y = np.where(sum_y > 0.5)[0]
  183. x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1]
  184. cv2.rectangle(img_array, (x0, y0), (x1, y1),
  185. tuple(color_mask.astype('int32').tolist()), 1)
  186. bbox_text = '%s %.2f' % (catid2name[catid], score)
  187. t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
  188. cv2.rectangle(img_array, (x0, y0), (x0 + t_size[0],
  189. y0 - t_size[1] - 3),
  190. tuple(color_mask.astype('int32').tolist()), -1)
  191. cv2.putText(
  192. img_array,
  193. bbox_text, (x0, y0 - 2),
  194. cv2.FONT_HERSHEY_SIMPLEX,
  195. 0.3, (0, 0, 0),
  196. 1,
  197. lineType=cv2.LINE_AA)
  198. return Image.fromarray(img_array.astype('uint8'))
  199. def map_coco_to_personlab(keypoints):
  200. permute = [0, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3]
  201. return keypoints[:, permute, :]
  202. def draw_pose(image, results, visual_thread=0.6, save_name='pose.jpg'):
  203. try:
  204. import matplotlib.pyplot as plt
  205. import matplotlib
  206. plt.switch_backend('agg')
  207. except Exception as e:
  208. logger.error('Matplotlib not found, plaese install matplotlib.'
  209. 'for example: `pip install matplotlib`.')
  210. raise e
  211. EDGES = [(0, 14), (0, 13), (0, 4), (0, 1), (14, 16), (13, 15), (4, 10),
  212. (1, 7), (10, 11), (7, 8), (11, 12), (8, 9), (4, 5), (1, 2),
  213. (5, 6), (2, 3)]
  214. NUM_EDGES = len(EDGES)
  215. colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
  216. [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
  217. [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
  218. cmap = matplotlib.cm.get_cmap('hsv')
  219. plt.figure()
  220. skeletons = np.array([item['keypoints'] for item in results]).reshape(
  221. -1, 17, 3)
  222. img = np.array(image).astype('float32')
  223. canvas = img.copy()
  224. for i in range(17):
  225. rgba = np.array(cmap(1 - i / 17. - 1. / 34))
  226. rgba[0:3] *= 255
  227. for j in range(len(skeletons)):
  228. if skeletons[j][i, 2] < visual_thread:
  229. continue
  230. cv2.circle(
  231. canvas,
  232. tuple(skeletons[j][i, 0:2].astype('int32')),
  233. 2,
  234. colors[i],
  235. thickness=-1)
  236. to_plot = cv2.addWeighted(img, 0.3, canvas, 0.7, 0)
  237. fig = matplotlib.pyplot.gcf()
  238. stickwidth = 2
  239. skeletons = map_coco_to_personlab(skeletons)
  240. for i in range(NUM_EDGES):
  241. for j in range(len(skeletons)):
  242. edge = EDGES[i]
  243. if skeletons[j][edge[0], 2] < visual_thread or skeletons[j][edge[
  244. 1], 2] < visual_thread:
  245. continue
  246. cur_canvas = canvas.copy()
  247. X = [skeletons[j][edge[0], 1], skeletons[j][edge[1], 1]]
  248. Y = [skeletons[j][edge[0], 0], skeletons[j][edge[1], 0]]
  249. mX = np.mean(X)
  250. mY = np.mean(Y)
  251. length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
  252. angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
  253. polygon = cv2.ellipse2Poly((int(mY), int(mX)),
  254. (int(length / 2), stickwidth),
  255. int(angle), 0, 360, 1)
  256. cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
  257. canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
  258. image = Image.fromarray(canvas.astype('uint8'))
  259. plt.close()
  260. return image