visualizer.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import cv2
  12. import numpy as np
  13. def get_color_map_list(length):
  14. """Returns the color map for visualizing the segmentation mask"""
  15. length += 1
  16. color_map = length * [0, 0, 0]
  17. for i in range(0, length):
  18. j = 0
  19. lab = i
  20. while lab:
  21. color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
  22. color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
  23. color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
  24. j += 1
  25. lab >>= 3
  26. color_map = color_map[3:]
  27. return color_map
  28. def visualize(image, result, weight=0.6, use_multilabel=False):
  29. """ Convert predict result to color image, and save added image. """
  30. color_map = get_color_map_list(256)
  31. color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
  32. color_map = np.array(color_map).astype("uint8")
  33. if not use_multilabel:
  34. # Use OpenCV LUT for color mapping
  35. c1 = cv2.LUT(result, color_map[:, 0])
  36. c2 = cv2.LUT(result, color_map[:, 1])
  37. c3 = cv2.LUT(result, color_map[:, 2])
  38. pseudo_img = np.dstack((c3, c2, c1))
  39. vis_result = cv2.addWeighted(image, weight, pseudo_img, 1 - weight, 0)
  40. else:
  41. vis_result = image.copy()
  42. for i in range(result.shape[0]):
  43. mask = result[i]
  44. c1 = np.where(mask, color_map[i, 0], vis_result[..., 0])
  45. c2 = np.where(mask, color_map[i, 1], vis_result[..., 1])
  46. c3 = np.where(mask, color_map[i, 2], vis_result[..., 2])
  47. pseudo_img = np.dstack((c3, c2, c1)).astype('uint8')
  48. contour, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
  49. cv2.CHAIN_APPROX_SIMPLE)
  50. vis_result = cv2.addWeighted(vis_result, weight, pseudo_img,
  51. 1 - weight, 0)
  52. contour_color = (int(color_map[i, 0]), int(color_map[i, 1]),
  53. int(color_map[i, 2]))
  54. vis_result = cv2.drawContours(vis_result, contour, -1,
  55. contour_color, 1)
  56. return vis_result