visualize.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. # Copyright (c) 2020 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. import os
  15. import cv2
  16. import numpy as np
  17. from PIL import Image as PILImage
  18. def visualize(image, result, save_dir=None, weight=0.6):
  19. """
  20. Convert predict result to color image, and save added image.
  21. Args:
  22. image (str): The path of origin image.
  23. result (np.ndarray): The predict result of image.
  24. save_dir (str): The directory for saving visual image. Default: None.
  25. weight (float): The image weight of visual image, and the result weight is (1 - weight). Default: 0.6
  26. Returns:
  27. vis_result (np.ndarray): If `save_dir` is None, return the visualized result.
  28. """
  29. color_map = get_color_map_list(256)
  30. color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
  31. color_map = np.array(color_map).astype("uint8")
  32. # Use OpenCV LUT for color mapping
  33. c1 = cv2.LUT(result, color_map[:, 0])
  34. c2 = cv2.LUT(result, color_map[:, 1])
  35. c3 = cv2.LUT(result, color_map[:, 2])
  36. pseudo_img = np.dstack((c1, c2, c3))
  37. im = cv2.imread(image)
  38. vis_result = cv2.addWeighted(im, weight, pseudo_img, 1 - weight, 0)
  39. if save_dir is not None:
  40. if not os.path.exists(save_dir):
  41. os.makedirs(save_dir)
  42. image_name = os.path.split(image)[-1]
  43. out_path = os.path.join(save_dir, image_name)
  44. cv2.imwrite(out_path, vis_result)
  45. else:
  46. return vis_result
  47. def get_pseudo_color_map(pred):
  48. pred_mask = PILImage.fromarray(pred.astype(np.uint8), mode='P')
  49. color_map = get_color_map_list(256)
  50. pred_mask.putpalette(color_map)
  51. return pred_mask
  52. def get_color_map_list(num_classes):
  53. """
  54. Returns the color map for visualizing the segmentation mask,
  55. which can support arbitrary number of classes.
  56. Args:
  57. num_classes (int): Number of classes.
  58. Returns:
  59. (list). The color map.
  60. """
  61. num_classes += 1
  62. color_map = num_classes * [0, 0, 0]
  63. for i in range(0, num_classes):
  64. j = 0
  65. lab = i
  66. while lab:
  67. color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
  68. color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
  69. color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
  70. j += 1
  71. lab >>= 3
  72. # color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
  73. color_map = color_map[3:]
  74. return color_map