topk.py 3.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. # copyright (c) 2021 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. import paddle
  17. import paddle.nn.functional as F
  18. class Topk(object):
  19. def __init__(self, topk=1, class_id_map_file=None):
  20. assert isinstance(topk, (int, ))
  21. self.class_id_map = self.parse_class_id_map(class_id_map_file)
  22. self.topk = topk
  23. def parse_class_id_map(self, class_id_map_file):
  24. if class_id_map_file is None:
  25. return None
  26. if not os.path.exists(class_id_map_file):
  27. print(
  28. "Warning: If want to use your own label_dict, please input legal path!\nOtherwise label_names will be empty!"
  29. )
  30. return None
  31. try:
  32. class_id_map = {}
  33. with open(class_id_map_file, "r") as fin:
  34. lines = fin.readlines()
  35. for line in lines:
  36. partition = line.split("\n")[0].partition(" ")
  37. class_id_map[int(partition[0])] = str(partition[-1])
  38. except Exception as ex:
  39. print(ex)
  40. class_id_map = None
  41. return class_id_map
  42. def __call__(self, x, file_names=None, multilabel=False):
  43. assert isinstance(x, paddle.Tensor)
  44. if file_names is not None:
  45. assert x.shape[0] == len(file_names)
  46. x = F.softmax(x, axis=-1) if not multilabel else F.sigmoid(x)
  47. x = x.numpy()
  48. y = []
  49. for idx, probs in enumerate(x):
  50. index = probs.argsort(axis=0)[-self.topk:][::-1].astype(
  51. "int32") if not multilabel else np.where(
  52. probs >= 0.5)[0].astype("int32")
  53. clas_id_list = []
  54. score_list = []
  55. label_name_list = []
  56. for i in index:
  57. clas_id_list.append(i.item())
  58. score_list.append(probs[i].item())
  59. if self.class_id_map is not None:
  60. label_name_list.append(self.class_id_map[i.item()])
  61. result = {
  62. "class_ids": clas_id_list,
  63. "scores": np.around(
  64. score_list, decimals=5).tolist(),
  65. }
  66. if file_names is not None:
  67. result["file_name"] = file_names[idx]
  68. if label_name_list is not None:
  69. result["label_names"] = label_name_list
  70. y.append(result)
  71. return y
  72. class MultiLabelTopk(Topk):
  73. def __init__(self, topk=1, class_id_map_file=None):
  74. super().__init__()
  75. def __call__(self, x, file_names=None):
  76. return super().__call__(x, file_names, multilabel=True)