clas.py 5.0 KB

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