clas.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  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 PIL
  15. from PIL import Image, ImageDraw, ImageFont
  16. import numpy as np
  17. import cv2
  18. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  19. from ..utils.color_map import get_colormap
  20. from .base import CVResult
  21. class TopkResult(CVResult):
  22. _HARD_FLAG = False
  23. def _to_img(self):
  24. """Draw label on image"""
  25. labels = self.get("label_names", self["class_ids"])
  26. label_str = f"{labels[0]} {self['scores'][0]:.2f}"
  27. image = self._img_reader.read(self["input_path"])
  28. if self._HARD_FLAG:
  29. image_np = np.array(image)
  30. image = Image.fromarray(image_np[:, :, ::-1])
  31. image_size = image.size
  32. draw = ImageDraw.Draw(image)
  33. min_font_size = int(image_size[0] * 0.02)
  34. max_font_size = int(image_size[0] * 0.05)
  35. for font_size in range(max_font_size, min_font_size - 1, -1):
  36. font = ImageFont.truetype(
  37. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8"
  38. )
  39. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  40. text_width_tmp, text_height_tmp = draw.textsize(label_str, font)
  41. else:
  42. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  43. text_width_tmp, text_height_tmp = right - left, bottom - top
  44. if text_width_tmp <= image_size[0]:
  45. break
  46. else:
  47. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, min_font_size)
  48. color_list = get_colormap(rgb=True)
  49. color = tuple(color_list[0])
  50. font_color = tuple(self._get_font_colormap(3))
  51. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  52. text_width, text_height = draw.textsize(label_str, font)
  53. else:
  54. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  55. text_width, text_height = right - left, bottom - top
  56. rect_left = 3
  57. rect_top = 3
  58. rect_right = rect_left + text_width + 3
  59. rect_bottom = rect_top + text_height + 6
  60. draw.rectangle([(rect_left, rect_top), (rect_right, rect_bottom)], fill=color)
  61. text_x = rect_left + 3
  62. text_y = rect_top
  63. draw.text((text_x, text_y), label_str, fill=font_color, font=font)
  64. return image
  65. def _get_font_colormap(self, color_index):
  66. """
  67. Get font colormap
  68. """
  69. dark = np.array([0x14, 0x0E, 0x35])
  70. light = np.array([0xFF, 0xFF, 0xFF])
  71. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  72. if color_index in light_indexs:
  73. return light.astype("int32")
  74. else:
  75. return dark.astype("int32")
  76. class MLClassResult(TopkResult):
  77. def _to_img(self):
  78. """Draw label on image"""
  79. image = self._img_reader.read(self["input_path"])
  80. label_names = self["label_names"]
  81. scores = self["scores"]
  82. image = image.convert("RGB")
  83. image_width, image_height = image.size
  84. font_size = int(image_width * 0.06)
  85. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size)
  86. text_lines = []
  87. row_width = 0
  88. row_height = 0
  89. row_text = "\t"
  90. for label_name, score in zip(label_names, scores):
  91. text = f"{label_name}({score})\t"
  92. if int(PIL.__version__.split(".")[0]) < 10:
  93. text_width, row_height = font.getsize(text)
  94. else:
  95. text_width, row_height = font.getbbox(text)[2:]
  96. if row_width + text_width <= image_width:
  97. row_text += text
  98. row_width += text_width
  99. else:
  100. text_lines.append(row_text)
  101. row_text = "\t" + text
  102. row_width = text_width
  103. text_lines.append(row_text)
  104. color_list = get_colormap(rgb=True)
  105. color = tuple(color_list[0])
  106. new_image_height = image_height + len(text_lines) * int(row_height * 1.2)
  107. new_image = Image.new("RGB", (image_width, new_image_height), color)
  108. new_image.paste(image, (0, 0))
  109. draw = ImageDraw.Draw(new_image)
  110. font_color = tuple(self._get_font_colormap(3))
  111. for i, text in enumerate(text_lines):
  112. if int(PIL.__version__.split(".")[0]) < 10:
  113. text_width, _ = font.getsize(text)
  114. else:
  115. text_width, _ = font.getbbox(text)[2:]
  116. draw.text(
  117. (0, image_height + i * int(row_height * 1.2)),
  118. text,
  119. fill=font_color,
  120. font=font,
  121. )
  122. return new_image