transforms.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  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 numpy as np
  16. from ...base import BaseTransform
  17. from .keys import ClsKeys as K
  18. from ....utils import logging
  19. __all__ = ["Topk", "NormalizeFeatures"]
  20. class Topk(BaseTransform):
  21. """ Topk Transform """
  22. def __init__(self, topk, class_id_map_file=None, delimiter=None):
  23. super().__init__()
  24. assert isinstance(topk, (int, ))
  25. self.topk = topk
  26. self.delimiter = delimiter if delimiter is not None else " "
  27. self.class_id_map = self._parse_class_id_map(class_id_map_file)
  28. def _parse_class_id_map(self, class_id_map_file):
  29. """ parse class id to label map file """
  30. if class_id_map_file is None:
  31. return None
  32. if not os.path.exists(class_id_map_file):
  33. logging.warning(
  34. "Warning: If want to use your own label_dict, please input legal path!\nOtherwise label_names will be empty!"
  35. )
  36. return None
  37. class_id_map = {}
  38. with open(class_id_map_file, 'r', encoding='utf-8') as fin:
  39. lines = fin.readlines()
  40. for line in lines:
  41. partition = line.split("\n")[0].partition(self.delimiter)
  42. class_id_map[int(partition[0])] = str(partition[-1])
  43. return class_id_map
  44. def apply(self, data):
  45. """ apply """
  46. x = data[K.CLS_PRED]
  47. class_id_map = data[
  48. K.
  49. LABELS] if self.class_id_map is None and K.LABELS in data else self.class_id_map
  50. y = []
  51. index = x.argsort(axis=0)[-self.topk:][::-1].astype("int32")
  52. clas_id_list = []
  53. score_list = []
  54. label_name_list = []
  55. for i in index:
  56. clas_id_list.append(i.item())
  57. score_list.append(x[i].item())
  58. if class_id_map is not None:
  59. label_name_list.append(class_id_map[i.item()])
  60. result = {
  61. "class_ids": clas_id_list,
  62. "scores": np.around(
  63. score_list, decimals=5).tolist()
  64. }
  65. if label_name_list is not None:
  66. result["label_names"] = label_name_list
  67. y.append(result)
  68. data[K.CLS_RESULT] = y
  69. return data
  70. @classmethod
  71. def get_input_keys(cls):
  72. """ get input keys """
  73. return [K.IM_PATH, K.CLS_PRED]
  74. @classmethod
  75. def get_output_keys(cls):
  76. """ get output keys """
  77. return [K.CLS_RESULT]
  78. class NormalizeFeatures(BaseTransform):
  79. """ Normalize Features Transform """
  80. def apply(self, data):
  81. """ apply """
  82. x = data[K.CLS_PRED]
  83. feas_norm = np.sqrt(np.sum(np.square(x), axis=0, keepdims=True))
  84. x = np.divide(x, feas_norm)
  85. data[K.CLS_RESULT] = x
  86. return data
  87. @classmethod
  88. def get_input_keys(cls):
  89. """ get input keys """
  90. return [K.IM_PATH, K.CLS_PRED]
  91. @classmethod
  92. def get_output_keys(cls):
  93. """ get output keys """
  94. return [K.CLS_RESULT]
  95. class PrintResult(BaseTransform):
  96. """ Print Result Transform """
  97. def apply(self, data):
  98. """ apply """
  99. logging.info("The prediction result is:")
  100. logging.info(data[K.CLS_RESULT])
  101. return data
  102. @classmethod
  103. def get_input_keys(cls):
  104. """ get input keys """
  105. return [K.CLS_RESULT]
  106. @classmethod
  107. def get_output_keys(cls):
  108. """ get output keys """
  109. return []
  110. class LoadLabels(BaseTransform):
  111. """load label to data
  112. """
  113. def __init__(self, labels=None):
  114. super().__init__()
  115. self.labels = labels
  116. def apply(self, data):
  117. """ apply """
  118. if self.labels:
  119. data[K.LABELS] = self.labels
  120. return data
  121. @classmethod
  122. def get_input_keys(cls):
  123. """ get input keys """
  124. return []
  125. @classmethod
  126. def get_output_keys(cls):
  127. """ get output keys """
  128. return [K.LABELS]