analyse_dataset.py 3.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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 platform
  16. from collections import defaultdict
  17. import matplotlib.pyplot as plt
  18. import numpy as np
  19. from matplotlib import font_manager
  20. from .....utils.file_interface import custom_open
  21. from .....utils.fonts import PINGFANG_FONT_FILE_PATH
  22. def deep_analyse(dataset_path, output):
  23. """class analysis for dataset"""
  24. tags = ["train", "val"]
  25. labels_cnt = defaultdict(str)
  26. label_path = os.path.join(dataset_path, "label.txt")
  27. with custom_open(label_path, "r") as f:
  28. lines = f.readlines()
  29. for line in lines:
  30. line = line.strip().split()
  31. labels_cnt[line[0]] = " ".join(line[1:])
  32. for tag in tags:
  33. anno_path = os.path.join(dataset_path, f"{tag}.txt")
  34. classes_num = defaultdict(int)
  35. for i in range(len(labels_cnt)):
  36. classes_num[labels_cnt[str(i)]] = 0
  37. with custom_open(anno_path, "r") as f:
  38. lines = f.readlines()
  39. for line in lines:
  40. line = line.strip().split()
  41. classes_num[labels_cnt[line[1]]] += 1
  42. if tag == "train":
  43. cnts_train = [cat_ids for cat_name, cat_ids in classes_num.items()]
  44. elif tag == "val":
  45. cnts_val = [cat_ids for cat_name, cat_ids in classes_num.items()]
  46. classes = [cat_name for cat_name, cat_ids in classes_num.items()]
  47. sorted_id = sorted(
  48. range(len(cnts_train)), key=lambda k: cnts_train[k], reverse=True
  49. )
  50. cnts_train_sorted = [cnts_train[index] for index in sorted_id]
  51. cnts_val_sorted = [cnts_val[index] for index in sorted_id]
  52. classes_sorted = [classes[index] for index in sorted_id]
  53. x = np.arange(len(classes))
  54. width = 0.5
  55. # bar
  56. os_system = platform.system().lower()
  57. if os_system == "windows":
  58. plt.rcParams["font.sans-serif"] = "FangSong"
  59. else:
  60. font = font_manager.FontProperties(fname=PINGFANG_FONT_FILE_PATH, size=10)
  61. fig, ax = plt.subplots(figsize=(max(8, int(len(classes) / 5)), 5), dpi=300)
  62. ax.bar(x, cnts_train_sorted, width=0.5, label="train")
  63. ax.bar(x + width, cnts_val_sorted, width=0.5, label="val")
  64. plt.xticks(
  65. x + width / 2,
  66. classes_sorted,
  67. rotation=90,
  68. fontproperties=None if os_system == "windows" else font,
  69. )
  70. ax.set_xlabel(
  71. "类别名称", fontproperties=None if os_system == "windows" else font, fontsize=12
  72. )
  73. ax.set_ylabel(
  74. "图片数量", fontproperties=None if os_system == "windows" else font, fontsize=12
  75. )
  76. plt.legend(loc=1)
  77. fig.tight_layout()
  78. file_path = os.path.join(output, "histogram.png")
  79. fig.savefig(file_path, dpi=300)
  80. return {"histogram": os.path.join("check_dataset", "histogram.png")}