transforms.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284
  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 os
  15. import json
  16. from pathlib import Path
  17. import numpy as np
  18. import PIL
  19. from PIL import ImageDraw, ImageFont
  20. from .keys import ClsKeys as K
  21. from ...base import BaseTransform
  22. from ...base.predictor.io import ImageWriter, ImageReader
  23. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  24. from ....utils import logging
  25. __all__ = ["Topk", "NormalizeFeatures", "PrintResult", "SaveClsResults"]
  26. def _parse_class_id_map(class_ids):
  27. """parse class id to label map file"""
  28. if class_ids is None:
  29. return None
  30. class_id_map = {id: str(lb) for id, lb in enumerate(class_ids)}
  31. return class_id_map
  32. class Topk(BaseTransform):
  33. """Topk Transform"""
  34. def __init__(self, topk, class_ids=None):
  35. super().__init__()
  36. assert isinstance(topk, (int,))
  37. self.topk = topk
  38. self.class_id_map = _parse_class_id_map(class_ids)
  39. def apply(self, data):
  40. """apply"""
  41. x = data[K.CLS_PRED]
  42. class_id_map = self.class_id_map
  43. y = []
  44. index = x.argsort(axis=0)[-self.topk :][::-1].astype("int32")
  45. clas_id_list = []
  46. score_list = []
  47. label_name_list = []
  48. for i in index:
  49. clas_id_list.append(i.item())
  50. score_list.append(x[i].item())
  51. if class_id_map is not None:
  52. label_name_list.append(class_id_map[i.item()])
  53. result = {
  54. "class_ids": clas_id_list,
  55. "scores": np.around(score_list, decimals=5).tolist(),
  56. }
  57. if label_name_list is not None:
  58. result["label_names"] = label_name_list
  59. y.append(result)
  60. data[K.CLS_RESULT] = y
  61. return data
  62. @classmethod
  63. def get_input_keys(cls):
  64. """get input keys"""
  65. return [K.IM_PATH, K.CLS_PRED]
  66. @classmethod
  67. def get_output_keys(cls):
  68. """get output keys"""
  69. return [K.CLS_RESULT]
  70. class NormalizeFeatures(BaseTransform):
  71. """Normalize Features Transform"""
  72. def apply(self, data):
  73. """apply"""
  74. x = data[K.CLS_PRED]
  75. feas_norm = np.sqrt(np.sum(np.square(x), axis=0, keepdims=True))
  76. x = np.divide(x, feas_norm)
  77. data[K.CLS_RESULT] = x
  78. return data
  79. @classmethod
  80. def get_input_keys(cls):
  81. """get input keys"""
  82. return [K.IM_PATH, K.CLS_PRED]
  83. @classmethod
  84. def get_output_keys(cls):
  85. """get output keys"""
  86. return [K.CLS_RESULT]
  87. class PrintResult(BaseTransform):
  88. """Print Result Transform"""
  89. def apply(self, data):
  90. """apply"""
  91. logging.info("The prediction result is:")
  92. logging.info(data[K.CLS_RESULT])
  93. return data
  94. @classmethod
  95. def get_input_keys(cls):
  96. """get input keys"""
  97. return [K.CLS_RESULT]
  98. @classmethod
  99. def get_output_keys(cls):
  100. """get output keys"""
  101. return []
  102. class SaveClsResults(BaseTransform):
  103. def __init__(self, save_dir, class_ids=None):
  104. super().__init__()
  105. self.save_dir = save_dir
  106. self.class_id_map = _parse_class_id_map(class_ids)
  107. self._writer = ImageWriter(backend="pillow")
  108. def _get_colormap(self, rgb=False):
  109. """
  110. Get colormap
  111. """
  112. color_list = np.array(
  113. [
  114. 0xFF,
  115. 0x00,
  116. 0x00,
  117. 0xCC,
  118. 0xFF,
  119. 0x00,
  120. 0x00,
  121. 0xFF,
  122. 0x66,
  123. 0x00,
  124. 0x66,
  125. 0xFF,
  126. 0xCC,
  127. 0x00,
  128. 0xFF,
  129. 0xFF,
  130. 0x4D,
  131. 0x00,
  132. 0x80,
  133. 0xFF,
  134. 0x00,
  135. 0x00,
  136. 0xFF,
  137. 0xB2,
  138. 0x00,
  139. 0x1A,
  140. 0xFF,
  141. 0xFF,
  142. 0x00,
  143. 0xE5,
  144. 0xFF,
  145. 0x99,
  146. 0x00,
  147. 0x33,
  148. 0xFF,
  149. 0x00,
  150. 0x00,
  151. 0xFF,
  152. 0xFF,
  153. 0x33,
  154. 0x00,
  155. 0xFF,
  156. 0xFF,
  157. 0x00,
  158. 0x99,
  159. 0xFF,
  160. 0xE5,
  161. 0x00,
  162. 0x00,
  163. 0xFF,
  164. 0x1A,
  165. 0x00,
  166. 0xB2,
  167. 0xFF,
  168. 0x80,
  169. 0x00,
  170. 0xFF,
  171. 0xFF,
  172. 0x00,
  173. 0x4D,
  174. ]
  175. ).astype(np.float32)
  176. color_list = color_list.reshape((-1, 3))
  177. if not rgb:
  178. color_list = color_list[:, ::-1]
  179. return color_list.astype("int32")
  180. def _get_font_colormap(self, color_index):
  181. """
  182. Get font colormap
  183. """
  184. dark = np.array([0x14, 0x0E, 0x35])
  185. light = np.array([0xFF, 0xFF, 0xFF])
  186. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  187. if color_index in light_indexs:
  188. return light.astype("int32")
  189. else:
  190. return dark.astype("int32")
  191. def apply(self, data):
  192. """Draw label on image"""
  193. ori_path = data[K.IM_PATH]
  194. pred = data[K.CLS_PRED]
  195. index = pred.argsort(axis=0)[-1].astype("int32")
  196. score = pred[index].item()
  197. label = self.class_id_map[int(index)] if self.class_id_map else ""
  198. label_str = f"{label} {score:.2f}"
  199. file_name = os.path.basename(ori_path)
  200. save_path = os.path.join(self.save_dir, file_name)
  201. image = ImageReader(backend="pil").read(ori_path)
  202. image = image.convert("RGB")
  203. image_size = image.size
  204. draw = ImageDraw.Draw(image)
  205. min_font_size = int(image_size[0] * 0.02)
  206. max_font_size = int(image_size[0] * 0.05)
  207. for font_size in range(max_font_size, min_font_size - 1, -1):
  208. font = ImageFont.truetype(
  209. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8"
  210. )
  211. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  212. text_width_tmp, text_height_tmp = draw.textsize(label_str, font)
  213. else:
  214. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  215. text_width_tmp, text_height_tmp = right - left, bottom - top
  216. if text_width_tmp <= image_size[0]:
  217. break
  218. else:
  219. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, min_font_size)
  220. color_list = self._get_colormap(rgb=True)
  221. color = tuple(color_list[0])
  222. font_color = tuple(self._get_font_colormap(3))
  223. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  224. text_width, text_height = draw.textsize(label_str, font)
  225. else:
  226. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  227. text_width, text_height = right - left, bottom - top
  228. rect_left = 3
  229. rect_top = 3
  230. rect_right = rect_left + text_width + 3
  231. rect_bottom = rect_top + text_height + 6
  232. draw.rectangle([(rect_left, rect_top), (rect_right, rect_bottom)], fill=color)
  233. text_x = rect_left + 3
  234. text_y = rect_top
  235. draw.text((text_x, text_y), label_str, fill=font_color, font=font)
  236. self._write_image(save_path, image)
  237. return data
  238. def _write_image(self, path, image):
  239. """write image"""
  240. if os.path.exists(path):
  241. logging.warning(f"{path} already exists. Overwriting it.")
  242. self._writer.write(path, image)
  243. @classmethod
  244. def get_input_keys(cls):
  245. """get input keys"""
  246. return [K.IM_PATH, K.CLS_PRED]
  247. @classmethod
  248. def get_output_keys(cls):
  249. """get output keys"""
  250. return []