seg.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  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 numpy as np
  15. import PIL
  16. from PIL import Image
  17. from .base import BaseResult
  18. class SegResult(BaseResult):
  19. """Save Result Transform"""
  20. def __init__(self, data):
  21. super().__init__(data)
  22. self.data = data
  23. # We use pillow backend to save both numpy arrays and PIL Image objects
  24. self._img_writer.set_backend("pillow", format_="PNG")
  25. def _get_res_img(self):
  26. """apply"""
  27. seg_map = self.data["pred"]
  28. pc_map = self.get_pseudo_color_map(seg_map[0])
  29. return pc_map
  30. def get_pseudo_color_map(self, pred):
  31. """get_pseudo_color_map"""
  32. if pred.min() < 0 or pred.max() > 255:
  33. raise ValueError("`pred` cannot be cast to uint8.")
  34. pred = pred.astype(np.uint8)
  35. pred_mask = Image.fromarray(pred, mode="P")
  36. color_map = self._get_color_map_list(256)
  37. pred_mask.putpalette(color_map)
  38. return pred_mask
  39. @staticmethod
  40. def _get_color_map_list(num_classes, custom_color=None):
  41. """_get_color_map_list"""
  42. num_classes += 1
  43. color_map = num_classes * [0, 0, 0]
  44. for i in range(0, num_classes):
  45. j = 0
  46. lab = i
  47. while lab:
  48. color_map[i * 3] |= ((lab >> 0) & 1) << (7 - j)
  49. color_map[i * 3 + 1] |= ((lab >> 1) & 1) << (7 - j)
  50. color_map[i * 3 + 2] |= ((lab >> 2) & 1) << (7 - j)
  51. j += 1
  52. lab >>= 3
  53. color_map = color_map[3:]
  54. if custom_color:
  55. color_map[: len(custom_color)] = custom_color
  56. return color_map