transforms.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import json
  13. from PIL import Image, ImageDraw, ImageFont
  14. from pathlib import Path
  15. import numpy as np
  16. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  17. from ...base import BaseTransform
  18. from ...base.predictor.io.writers import ImageWriter
  19. from .keys import ClsKeys as K
  20. from ....utils import logging
  21. __all__ = ["Topk", "NormalizeFeatures", "PrintResult", "SaveClsResults"]
  22. def _parse_class_id_map(class_ids):
  23. """ parse class id to label map file """
  24. if class_ids is None:
  25. return None
  26. class_id_map = {id: str(lb) for id, lb in enumerate(class_ids)}
  27. return class_id_map
  28. class Topk(BaseTransform):
  29. """ Topk Transform """
  30. def __init__(self, topk, class_ids=None):
  31. super().__init__()
  32. assert isinstance(topk, (int, ))
  33. self.topk = topk
  34. self.class_id_map = _parse_class_id_map(class_ids)
  35. def apply(self, data):
  36. """ apply """
  37. x = data[K.CLS_PRED]
  38. class_id_map = self.class_id_map
  39. y = []
  40. index = x.argsort(axis=0)[-self.topk:][::-1].astype("int32")
  41. clas_id_list = []
  42. score_list = []
  43. label_name_list = []
  44. for i in index:
  45. clas_id_list.append(i.item())
  46. score_list.append(x[i].item())
  47. if class_id_map is not None:
  48. label_name_list.append(class_id_map[i.item()])
  49. result = {
  50. "class_ids": clas_id_list,
  51. "scores": np.around(
  52. score_list, decimals=5).tolist()
  53. }
  54. if label_name_list is not None:
  55. result["label_names"] = label_name_list
  56. y.append(result)
  57. data[K.CLS_RESULT] = y
  58. return data
  59. @classmethod
  60. def get_input_keys(cls):
  61. """ get input keys """
  62. return [K.IM_PATH, K.CLS_PRED]
  63. @classmethod
  64. def get_output_keys(cls):
  65. """ get output keys """
  66. return [K.CLS_RESULT]
  67. class NormalizeFeatures(BaseTransform):
  68. """ Normalize Features Transform """
  69. def apply(self, data):
  70. """ apply """
  71. x = data[K.CLS_PRED]
  72. feas_norm = np.sqrt(np.sum(np.square(x), axis=0, keepdims=True))
  73. x = np.divide(x, feas_norm)
  74. data[K.CLS_RESULT] = x
  75. return data
  76. @classmethod
  77. def get_input_keys(cls):
  78. """ get input keys """
  79. return [K.IM_PATH, K.CLS_PRED]
  80. @classmethod
  81. def get_output_keys(cls):
  82. """ get output keys """
  83. return [K.CLS_RESULT]
  84. class PrintResult(BaseTransform):
  85. """ Print Result Transform """
  86. def apply(self, data):
  87. """ apply """
  88. logging.info("The prediction result is:")
  89. logging.info(data[K.CLS_RESULT])
  90. return data
  91. @classmethod
  92. def get_input_keys(cls):
  93. """ get input keys """
  94. return [K.CLS_RESULT]
  95. @classmethod
  96. def get_output_keys(cls):
  97. """ get output keys """
  98. return []
  99. class SaveClsResults(BaseTransform):
  100. def __init__(self, save_dir, class_ids=None):
  101. super().__init__()
  102. self.save_dir = save_dir
  103. self.class_id_map = _parse_class_id_map(class_ids)
  104. self._writer = ImageWriter(backend='pillow')
  105. def _get_colormap(self, rgb=False):
  106. """
  107. Get colormap
  108. """
  109. color_list = np.array([
  110. 0xFF, 0x00, 0x00, 0xCC, 0xFF, 0x00, 0x00, 0xFF, 0x66, 0x00, 0x66,
  111. 0xFF, 0xCC, 0x00, 0xFF, 0xFF, 0x4D, 0x00, 0x80, 0xff, 0x00, 0x00,
  112. 0xFF, 0xB2, 0x00, 0x1A, 0xFF, 0xFF, 0x00, 0xE5, 0xFF, 0x99, 0x00,
  113. 0x33, 0xFF, 0x00, 0x00, 0xFF, 0xFF, 0x33, 0x00, 0xFF, 0xff, 0x00,
  114. 0x99, 0xFF, 0xE5, 0x00, 0x00, 0xFF, 0x1A, 0x00, 0xB2, 0xFF, 0x80,
  115. 0x00, 0xFF, 0xFF, 0x00, 0x4D
  116. ]).astype(np.float32)
  117. color_list = (color_list.reshape((-1, 3)))
  118. if not rgb:
  119. color_list = color_list[:, ::-1]
  120. return color_list.astype('int32')
  121. def _get_font_colormap(self, color_index):
  122. """
  123. Get font colormap
  124. """
  125. dark = np.array([0x14, 0x0E, 0x35])
  126. light = np.array([0xFF, 0xFF, 0xFF])
  127. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  128. if color_index in light_indexs:
  129. return light.astype('int32')
  130. else:
  131. return dark.astype('int32')
  132. def apply(self, data):
  133. """ Draw label on image """
  134. ori_path = data[K.IM_PATH]
  135. pred = data[K.CLS_PRED]
  136. index = pred.argsort(axis=0)[-1].astype("int32")
  137. score = pred[index].item()
  138. label = self.class_id_map[int(index)] if self.class_id_map else ""
  139. label_str = f"{label} {score:.2f}"
  140. file_name = os.path.basename(ori_path)
  141. save_path = os.path.join(self.save_dir, file_name)
  142. image = Image.open(ori_path)
  143. image = image.convert('RGB')
  144. image_size = image.size
  145. draw = ImageDraw.Draw(image)
  146. min_font_size = int(image_size[0] * 0.02)
  147. max_font_size = int(image_size[0] * 0.05)
  148. for font_size in range(max_font_size, min_font_size - 1, -1):
  149. font = ImageFont.truetype(
  150. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  151. text_width_tmp, text_height_tmp = draw.textsize(label_str, font)
  152. if text_width_tmp <= image_size[0]:
  153. break
  154. else:
  155. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH,
  156. min_font_size)
  157. color_list = self._get_colormap(rgb=True)
  158. color = tuple(color_list[0])
  159. font_color = tuple(self._get_font_colormap(3))
  160. text_width, text_height = draw.textsize(label_str, font)
  161. rect_left = 3
  162. rect_top = 3
  163. rect_right = rect_left + text_width + 3
  164. rect_bottom = rect_top + text_height + 6
  165. draw.rectangle(
  166. [(rect_left, rect_top), (rect_right, rect_bottom)], fill=color)
  167. text_x = rect_left + 3
  168. text_y = rect_top
  169. draw.text((text_x, text_y), label_str, fill=font_color, font=font)
  170. self._write_image(save_path, image)
  171. return data
  172. def _write_image(self, path, image):
  173. """ write image """
  174. if os.path.exists(path):
  175. logging.warning(f"{path} already exists. Overwriting it.")
  176. self._writer.write(path, image)
  177. @classmethod
  178. def get_input_keys(cls):
  179. """ get input keys """
  180. return [K.IM_PATH, K.CLS_PRED]
  181. @classmethod
  182. def get_output_keys(cls):
  183. """ get output keys """
  184. return []