clas.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  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. image = image.convert("RGB")
  77. label_names = self["label_names"]
  78. scores = self["scores"]
  79. image = image.convert("RGB")
  80. image_width, image_height = image.size
  81. font_size = int(image_width * 0.06)
  82. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size)
  83. text_lines = []
  84. row_width = 0
  85. row_height = 0
  86. row_text = "\t"
  87. for label_name, score in zip(label_names, scores):
  88. text = f"{label_name}({score})\t"
  89. text_width, row_height = font.getsize(text)
  90. if row_width + text_width <= image_width:
  91. row_text += text
  92. row_width += text_width
  93. else:
  94. text_lines.append(row_text)
  95. row_text = "\t" + text
  96. row_width = text_width
  97. text_lines.append(row_text)
  98. color_list = get_colormap(rgb=True)
  99. color = tuple(color_list[0])
  100. new_image_height = image_height + len(text_lines) * int(row_height * 1.2)
  101. new_image = Image.new("RGB", (image_width, new_image_height), color)
  102. new_image.paste(image, (0, 0))
  103. draw = ImageDraw.Draw(new_image)
  104. font_color = tuple(self._get_font_colormap(3))
  105. for i, text in enumerate(text_lines):
  106. text_width, _ = font.getsize(text)
  107. draw.text(
  108. (0, image_height + i * int(row_height * 1.2)),
  109. text,
  110. fill=font_color,
  111. font=font,
  112. )
  113. return new_image