transforms.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. # copyright (c) 2024 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 numpy as np
  16. from PIL import Image
  17. from ....utils import logging
  18. from ...base import BaseTransform
  19. from ...base.predictor.io.writers import ImageWriter
  20. from .keys import SegKeys as K
  21. from skimage import measure, morphology
  22. __all__ = ["GeneratePCMap", "SaveSegResults"]
  23. class GeneratePCMap(BaseTransform):
  24. """GeneratePCMap"""
  25. def __init__(self, color_map=None):
  26. super().__init__()
  27. self.color_map = color_map
  28. def apply(self, data):
  29. """apply"""
  30. pred = data[K.SEG_MAP]
  31. pc_map = self.get_pseudo_color_map(pred)
  32. data[K.PC_MAP] = pc_map
  33. return data
  34. @classmethod
  35. def get_input_keys(cls):
  36. """get input keys"""
  37. return [K.SEG_MAP]
  38. @classmethod
  39. def get_output_keys(cls):
  40. """get input keys"""
  41. return [K.PC_MAP]
  42. def get_pseudo_color_map(self, pred):
  43. """get_pseudo_color_map"""
  44. if pred.min() < 0 or pred.max() > 255:
  45. raise ValueError("`pred` cannot be cast to uint8.")
  46. pred = pred.astype(np.uint8)
  47. pred_mask = Image.fromarray(pred, mode="P")
  48. if self.color_map is None:
  49. color_map = self._get_color_map_list(256)
  50. else:
  51. color_map = self.color_map
  52. pred_mask.putpalette(color_map)
  53. return pred_mask
  54. @staticmethod
  55. def _get_color_map_list(num_classes, custom_color=None):
  56. """_get_color_map_list"""
  57. num_classes += 1
  58. color_map = num_classes * [0, 0, 0]
  59. for i in range(0, num_classes):
  60. j = 0
  61. lab = i
  62. while lab:
  63. color_map[i * 3] |= ((lab >> 0) & 1) << (7 - j)
  64. color_map[i * 3 + 1] |= ((lab >> 1) & 1) << (7 - j)
  65. color_map[i * 3 + 2] |= ((lab >> 2) & 1) << (7 - j)
  66. j += 1
  67. lab >>= 3
  68. color_map = color_map[3:]
  69. if custom_color:
  70. color_map[: len(custom_color)] = custom_color
  71. return color_map
  72. class SaveSegResults(BaseTransform):
  73. """SaveSegResults"""
  74. _PC_MAP_SUFFIX = "_pc"
  75. _FILE_EXT = ".png"
  76. def __init__(self, save_dir, save_pc_map=True):
  77. super().__init__()
  78. self.save_dir = save_dir
  79. self.save_pc_map = save_pc_map
  80. # We use pillow backend to save both numpy arrays and PIL Image objects
  81. self._writer = ImageWriter(backend="pillow")
  82. def apply(self, data):
  83. """apply"""
  84. seg_map = data[K.SEG_MAP]
  85. ori_path = data[K.IM_PATH]
  86. file_name = os.path.basename(ori_path)
  87. file_name = self._replace_ext(file_name, self._FILE_EXT)
  88. seg_map_save_path = os.path.join(self.save_dir, file_name)
  89. self._write_im(seg_map_save_path, seg_map)
  90. if self.save_pc_map:
  91. if K.PC_MAP in data:
  92. pc_map_save_path = self._add_suffix(
  93. seg_map_save_path, self._PC_MAP_SUFFIX
  94. )
  95. pc_map = data[K.PC_MAP]
  96. self._write_im(pc_map_save_path, pc_map)
  97. else:
  98. logging.warning(f"The {K.PC_MAP} result don't exist!")
  99. return data
  100. @classmethod
  101. def get_input_keys(cls):
  102. """get input keys"""
  103. return [K.IM_PATH, K.SEG_MAP]
  104. @classmethod
  105. def get_output_keys(cls):
  106. """get output keys"""
  107. return []
  108. def _write_im(self, path, im):
  109. """write image"""
  110. if os.path.exists(path):
  111. logging.warning(f"{path} already exists. Overwriting it.")
  112. self._writer.write(path, im)
  113. @staticmethod
  114. def _add_suffix(path, suffix):
  115. """add suffix"""
  116. stem, ext = os.path.splitext(path)
  117. return stem + suffix + ext
  118. @staticmethod
  119. def _replace_ext(path, new_ext):
  120. """replace ext"""
  121. stem, _ = os.path.splitext(path)
  122. return stem + new_ext
  123. class PrintResult(BaseTransform):
  124. """Print Result Transform"""
  125. def apply(self, data):
  126. """apply"""
  127. logging.info("The prediction result is:")
  128. logging.info(f"keys: {data.keys()}")
  129. return data
  130. @classmethod
  131. def get_input_keys(cls):
  132. """get input keys"""
  133. return [K.SEG_MAP]
  134. @classmethod
  135. def get_output_keys(cls):
  136. """get output keys"""
  137. return []
  138. class Map_to_mask(BaseTransform):
  139. """Map_to_mask"""
  140. def apply(self, data):
  141. """apply"""
  142. score_map = data[K.SEG_MAP]
  143. thred = 0.01
  144. mask = score_map[0]
  145. mask[mask > thred] = 255
  146. mask[mask <= thred] = 0
  147. kernel = morphology.disk(4)
  148. mask = morphology.opening(mask, kernel)
  149. mask = mask.astype(np.uint8)
  150. data[K.SEG_MAP] = mask
  151. return data
  152. @classmethod
  153. def get_input_keys(cls):
  154. """get input keys"""
  155. return [K.SEG_MAP]
  156. @classmethod
  157. def get_output_keys(cls):
  158. """get output keys"""
  159. return []