analyse_dataset.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  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 math
  15. import os
  16. import platform
  17. from collections import defaultdict
  18. import cv2
  19. import matplotlib.pyplot as plt
  20. import numpy as np
  21. from matplotlib import font_manager
  22. from matplotlib.backends.backend_agg import FigureCanvasAgg
  23. from .....utils.file_interface import custom_open
  24. from .....utils.fonts import PINGFANG_FONT_FILE_PATH
  25. from .....utils.logging import warning
  26. def simple_analyse(dataset_path, images_dict):
  27. """
  28. Analyse the dataset samples by return image path and label path
  29. Args:
  30. dataset_path (str): dataset path
  31. ds_meta (dict): dataset meta
  32. images_dict (dict): train, val and test image path
  33. Returns:
  34. tuple: tuple of sample number, image path and label path for train, val and text subdataset.
  35. """
  36. tags = ["train", "val", "test"]
  37. sample_cnts = defaultdict(int)
  38. img_paths = defaultdict(list)
  39. res = [None] * 6
  40. for tag in tags:
  41. file_list = os.path.join(dataset_path, f"{tag}.txt")
  42. if not os.path.exists(file_list):
  43. if tag in ("train", "val"):
  44. res.insert(0, "数据集不符合规范,请先通过数据校准")
  45. return res
  46. else:
  47. continue
  48. else:
  49. with custom_open(file_list, "r") as f:
  50. all_lines = f.readlines()
  51. # Each line corresponds to a sample
  52. sample_cnts[tag] = len(all_lines)
  53. img_paths[tag] = images_dict[tag]
  54. return (
  55. "完成数据分析",
  56. sample_cnts[tags[0]],
  57. sample_cnts[tags[1]],
  58. sample_cnts[tags[2]],
  59. img_paths[tags[0]],
  60. img_paths[tags[1]],
  61. img_paths[tags[2]],
  62. )
  63. def deep_analyse(dataset_path, output, datatype="MSTextRecDataset"):
  64. """class analysis for dataset"""
  65. tags = ["train", "val"]
  66. labels_cnt = {}
  67. x_max = []
  68. classes_max = []
  69. for tag in tags:
  70. image_path = os.path.join(dataset_path, f"{tag}.txt")
  71. str_nums = []
  72. with custom_open(image_path, "r") as f:
  73. lines = f.readlines()
  74. for line in lines:
  75. line = line.strip().split("\t")
  76. if len(line) != 2:
  77. warning(f"Error in {line}.")
  78. continue
  79. str_nums.append(len(line[1]))
  80. if datatype == "LaTeXOCRDataset":
  81. max_length = min(768, max(str_nums))
  82. interval = 20
  83. else:
  84. max_length = min(100, max(str_nums))
  85. interval = 5
  86. start = 0
  87. for i in range(1, math.ceil((max_length / interval))):
  88. stop = i * interval
  89. num_str = sum(start < i <= stop for i in str_nums)
  90. labels_cnt[f"{start}-{stop}"] = num_str
  91. start = stop
  92. if sum(max_length < i for i in str_nums) != 0:
  93. labels_cnt[f"> {max_length}"] = sum(max_length < i for i in str_nums)
  94. if tag == "train":
  95. cnts_train = [cat_ids for cat_name, cat_ids in labels_cnt.items()]
  96. x_train = np.arange(len(cnts_train))
  97. if len(x_train) > len(x_max):
  98. x_max = x_train
  99. classes_max = [cat_name for cat_name, cat_ids in labels_cnt.items()]
  100. elif tag == "val":
  101. cnts_val = [cat_ids for cat_name, cat_ids in labels_cnt.items()]
  102. x_val = np.arange(len(cnts_val))
  103. if len(x_val) > len(x_max):
  104. x_max = x_val
  105. classes_max = [cat_name for cat_name, cat_ids in labels_cnt.items()]
  106. width = 0.3
  107. # bar
  108. os_system = platform.system().lower()
  109. if os_system == "windows":
  110. plt.rcParams["font.sans-serif"] = "FangSong"
  111. else:
  112. font = font_manager.FontProperties(fname=PINGFANG_FONT_FILE_PATH, size=15)
  113. if datatype == "LaTeXOCRDataset":
  114. fig, ax = plt.subplots(figsize=(15, 9), dpi=120)
  115. xlabel_name = "公式长度区间"
  116. else:
  117. fig, ax = plt.subplots(figsize=(10, 5), dpi=120)
  118. xlabel_name = "文本字长度区间"
  119. ax.bar(x_train, cnts_train, width=0.3, label="train")
  120. ax.bar(x_val + width, cnts_val, width=0.3, label="val")
  121. plt.xticks(x_max + width / 2, classes_max, rotation=90)
  122. plt.legend(prop={"size": 18})
  123. ax.set_xlabel(
  124. xlabel_name,
  125. fontproperties=None if os_system == "windows" else font,
  126. fontsize=12,
  127. )
  128. ax.set_ylabel(
  129. "图片数量", fontproperties=None if os_system == "windows" else font, fontsize=12
  130. )
  131. canvas = FigureCanvasAgg(fig)
  132. canvas.draw()
  133. width, height = fig.get_size_inches() * fig.get_dpi()
  134. pie_array = np.frombuffer(canvas.tostring_rgb(), dtype="uint8").reshape(
  135. int(height), int(width), 3
  136. )
  137. fig1_path = os.path.join(output, "histogram.png")
  138. cv2.imwrite(fig1_path, pie_array)
  139. return {"histogram": os.path.join("check_dataset", "histogram.png")}