visualize.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. #copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  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 os
  15. import cv2
  16. import numpy as np
  17. from PIL import Image, ImageDraw
  18. def visualize_detection(image, result, threshold=0.5, save_dir=None):
  19. """
  20. Visualize bbox and mask results
  21. """
  22. image_name = os.path.split(image)[-1]
  23. image = Image.open(image).convert('RGB')
  24. image = draw_bbox_mask(image, result, threshold=threshold)
  25. if save_dir is not None:
  26. if not os.path.exists(save_dir):
  27. os.makedirs(save_dir)
  28. out_path = os.path.join(save_dir, 'visualize_{}'.format(image_name))
  29. image.save(out_path, quality=95)
  30. else:
  31. return image
  32. def visualize_segmentation(image, result, weight=0.6, save_dir=None):
  33. """
  34. Convert segment result to color image, and save added image.
  35. Args:
  36. image: the path of origin image
  37. result: the predict result of image
  38. weight: the image weight of visual image, and the result weight is (1 - weight)
  39. save_dir: the directory for saving visual image
  40. """
  41. label_map = result['label_map']
  42. color_map = get_color_map_list(256)
  43. color_map = np.array(color_map).astype("uint8")
  44. # Use OpenCV LUT for color mapping
  45. c1 = cv2.LUT(label_map, color_map[:, 0])
  46. c2 = cv2.LUT(label_map, color_map[:, 1])
  47. c3 = cv2.LUT(label_map, color_map[:, 2])
  48. pseudo_img = np.dstack((c1, c2, c3))
  49. im = cv2.imread(image)
  50. vis_result = cv2.addWeighted(im, weight, pseudo_img, 1 - weight, 0)
  51. if save_dir is not None:
  52. if not os.path.exists(save_dir):
  53. os.makedirs(save_dir)
  54. image_name = os.path.split(image)[-1]
  55. out_path = os.path.join(save_dir, 'visualize_{}'.format(image_name))
  56. cv2.imwrite(out_path, vis_result)
  57. else:
  58. return vis_result
  59. def get_color_map_list(num_classes):
  60. """ Returns the color map for visualizing the segmentation mask,
  61. which can support arbitrary number of classes.
  62. Args:
  63. num_classes: Number of classes
  64. Returns:
  65. The color map
  66. """
  67. color_map = num_classes * [0, 0, 0]
  68. for i in range(0, num_classes):
  69. j = 0
  70. lab = i
  71. while lab:
  72. color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
  73. color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
  74. color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
  75. j += 1
  76. lab >>= 3
  77. color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
  78. return color_map
  79. # expand an array of boxes by a given scale.
  80. def expand_boxes(boxes, scale):
  81. """
  82. """
  83. w_half = (boxes[:, 2] - boxes[:, 0]) * .5
  84. h_half = (boxes[:, 3] - boxes[:, 1]) * .5
  85. x_c = (boxes[:, 2] + boxes[:, 0]) * .5
  86. y_c = (boxes[:, 3] + boxes[:, 1]) * .5
  87. w_half *= scale
  88. h_half *= scale
  89. boxes_exp = np.zeros(boxes.shape)
  90. boxes_exp[:, 0] = x_c - w_half
  91. boxes_exp[:, 2] = x_c + w_half
  92. boxes_exp[:, 1] = y_c - h_half
  93. boxes_exp[:, 3] = y_c + h_half
  94. return boxes_exp
  95. def clip_bbox(bbox):
  96. xmin = max(min(bbox[0], 1.), 0.)
  97. ymin = max(min(bbox[1], 1.), 0.)
  98. xmax = max(min(bbox[2], 1.), 0.)
  99. ymax = max(min(bbox[3], 1.), 0.)
  100. return xmin, ymin, xmax, ymax
  101. def draw_bbox_mask(image, results, threshold=0.5, alpha=0.7):
  102. labels = list()
  103. for dt in np.array(results):
  104. if dt['category'] not in labels:
  105. labels.append(dt['category'])
  106. color_map = get_color_map_list(len(labels))
  107. for dt in np.array(results):
  108. cname, bbox, score = dt['category'], dt['bbox'], dt['score']
  109. if score < threshold:
  110. continue
  111. xmin, ymin, w, h = bbox
  112. xmax = xmin + w
  113. ymax = ymin + h
  114. color = tuple(color_map[labels.index(cname)])
  115. # draw bbox
  116. draw = ImageDraw.Draw(image)
  117. draw.line([(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
  118. (xmin, ymin)],
  119. width=2,
  120. fill=color)
  121. # draw label
  122. text = "{} {:.2f}".format(cname, score)
  123. tw, th = draw.textsize(text)
  124. draw.rectangle([(xmin + 1, ymin - th), (xmin + tw + 1, ymin)],
  125. fill=color)
  126. draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
  127. # draw mask
  128. if 'mask' in dt:
  129. mask = dt['mask']
  130. color_mask = np.array(color_map[labels.index(
  131. dt['category'])]).astype('float32')
  132. img_array = np.array(image).astype('float32')
  133. idx = np.nonzero(mask)
  134. img_array[idx[0], idx[1], :] *= 1.0 - alpha
  135. img_array[idx[0], idx[1], :] += alpha * color_mask
  136. image = Image.fromarray(img_array.astype('uint8'))
  137. return image