analyse_dataset.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  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. import matplotlib.pyplot as plt
  17. import numpy as np
  18. from matplotlib import font_manager
  19. from .....utils.file_interface import custom_open
  20. from .....utils.fonts import PINGFANG_FONT_FILE_PATH
  21. def deep_analyse(dataset_path, output, dataset_type="ShiTuRec"):
  22. """class analysis for dataset"""
  23. tags = ["train", "gallery", "query"]
  24. tags_info = dict()
  25. for tag in tags:
  26. anno_path = os.path.join(dataset_path, f"{tag}.txt")
  27. with custom_open(anno_path, "r") as f:
  28. lines = f.readlines()
  29. lines = [line.strip("\n").split(" ") for line in lines]
  30. num_images = len(lines)
  31. num_labels = len(set([int(line[1]) for line in lines]))
  32. tags_info[tag] = {
  33. "num_images": num_images,
  34. "num_labels": num_labels,
  35. }
  36. categories = list(tags_info.keys())
  37. num_images = [tags_info[category]["num_images"] for category in categories]
  38. num_labels = [tags_info[category]["num_labels"] for category in categories]
  39. # bar
  40. os_system = platform.system().lower()
  41. if os_system == "windows":
  42. plt.rcParams["font.sans-serif"] = "FangSong"
  43. else:
  44. font = font_manager.FontProperties(fname=PINGFANG_FONT_FILE_PATH, size=10)
  45. x = np.arange(len(categories)) # 标签位置
  46. width = 0.35 # 每个条形的宽度
  47. fig, ax = plt.subplots()
  48. rects1 = ax.bar(x - width / 2, num_images, width, label="Num Images")
  49. rects2 = ax.bar(x + width / 2, num_labels, width, label="Num Classes")
  50. # 添加一些文本标签
  51. ax.set_xlabel("集合", fontproperties=None if os_system == "windows" else font)
  52. ax.set_ylabel("数量", fontproperties=None if os_system == "windows" else font)
  53. ax.set_title(
  54. "不同集合的图片和类别数量",
  55. fontproperties=None if os_system == "windows" else font,
  56. )
  57. ax.set_xticks(x, fontproperties=None if os_system == "windows" else font)
  58. ax.set_xticklabels(categories)
  59. ax.legend()
  60. # 在条形图上添加数值标签
  61. def autolabel(rects):
  62. """Attach a text label above each bar in *rects*, displaying its height."""
  63. for rect in rects:
  64. height = rect.get_height()
  65. ax.annotate(
  66. "{}".format(height),
  67. xy=(rect.get_x() + rect.get_width() / 2, height),
  68. xytext=(0, 3), # 3 points vertical offset
  69. textcoords="offset points",
  70. ha="center",
  71. va="bottom",
  72. )
  73. autolabel(rects1)
  74. autolabel(rects2)
  75. fig.tight_layout()
  76. file_path = os.path.join(output, "histogram.png")
  77. fig.savefig(file_path, dpi=300)
  78. return {"histogram": os.path.join("check_dataset", "histogram.png")}