analyse_dataset.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  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 json
  15. import os
  16. from collections import defaultdict
  17. import cv2
  18. import matplotlib.pyplot as plt
  19. import numpy as np
  20. from matplotlib.backends.backend_agg import FigureCanvasAgg
  21. from PIL import Image, ImageOps
  22. from .....utils.file_interface import custom_open
  23. # show data samples
  24. def simple_analyse(dataset_path, max_recorded_sample_cnts=20, show_label=True):
  25. """
  26. Analyse the dataset samples by return not nore than
  27. max_recorded_sample_cnts image path and label path
  28. Args:
  29. dataset_path (str): dataset path
  30. max_recorded_sample_cnts (int, optional): the number to return. Default: 50.
  31. Returns:
  32. tuple: tuple of sample number, image path and label path for train, val and text subdataset.
  33. """
  34. tags = ["train", "val", "test"]
  35. sample_cnts = defaultdict(int)
  36. img_paths = defaultdict(list)
  37. lab_paths = defaultdict(list)
  38. lab_infos = defaultdict(list)
  39. res = [None] * 9
  40. delim = "\t"
  41. for tag in tags:
  42. file_list = os.path.join(dataset_path, f"{tag}.txt")
  43. if not os.path.exists(file_list):
  44. if tag in ("train", "val"):
  45. res.insert(0, "数据集不符合规范,请先通过数据校准")
  46. return res
  47. else:
  48. continue
  49. else:
  50. with custom_open(file_list, "r") as f:
  51. all_lines = f.readlines()
  52. # Each line corresponds to a sample
  53. sample_cnts[tag] = len(all_lines)
  54. for idx, line in enumerate(all_lines):
  55. parts = line.strip("\n").split(delim)
  56. if len(line.strip("\n")) < 1:
  57. continue
  58. if tag in ("train", "val"):
  59. valid_num_parts_lst = [2]
  60. else:
  61. valid_num_parts_lst = [1, 2]
  62. if len(parts) not in valid_num_parts_lst and len(line.strip("\n")) > 1:
  63. res.insert(0, "数据集的标注文件不符合规范")
  64. return res
  65. if len(parts) == 2:
  66. img_path, lab_path = parts
  67. else:
  68. # len(parts) == 1
  69. img_path = parts[0]
  70. lab_path = None
  71. # check det label
  72. if len(img_paths[tag]) < max_recorded_sample_cnts:
  73. img_path = os.path.join(dataset_path, img_path)
  74. if lab_path is not None:
  75. label = json.loads(lab_path)
  76. boxes = []
  77. for item in label:
  78. if "points" not in item or "transcription" not in item:
  79. res.insert(0, "数据集的标注文件不符合规范")
  80. return res
  81. box = np.array(item["points"])
  82. if box.shape[1] != 2:
  83. res.insert(0, "数据集的标注文件不符合规范")
  84. return res
  85. boxes.append(box)
  86. txt = item["transcription"]
  87. if not isinstance(txt, str):
  88. res.insert(0, "数据集的标注文件不符合规范")
  89. return res
  90. if show_label:
  91. lab_img = show_label_img(img_path, boxes)
  92. img_paths[tag].append(img_path)
  93. if show_label:
  94. lab_paths[tag].append(lab_img)
  95. else:
  96. lab_infos[tag].append({"img_path": img_path, "box": boxes})
  97. if show_label:
  98. return (
  99. "完成数据分析",
  100. sample_cnts[tags[0]],
  101. sample_cnts[tags[1]],
  102. sample_cnts[tags[2]],
  103. img_paths[tags[0]],
  104. img_paths[tags[1]],
  105. img_paths[tags[2]],
  106. lab_paths[tags[0]],
  107. lab_paths[tags[1]],
  108. lab_paths[tags[2]],
  109. )
  110. else:
  111. return (
  112. "完成数据分析",
  113. sample_cnts[tags[0]],
  114. sample_cnts[tags[1]],
  115. sample_cnts[tags[2]],
  116. img_paths[tags[0]],
  117. img_paths[tags[1]],
  118. img_paths[tags[2]],
  119. lab_infos[tags[0]],
  120. lab_infos[tags[1]],
  121. lab_infos[tags[2]],
  122. )
  123. def show_label_img(img_path, dt_boxes):
  124. """draw ocr detection label"""
  125. img = cv2.imread(img_path)
  126. for box in dt_boxes:
  127. box = np.array(box).astype(np.int32).reshape(-1, 2)
  128. cv2.polylines(img, [box], True, color=(0, 255, 0), thickness=3)
  129. return img[:, :, ::-1]
  130. def deep_analyse(dataset_path, output):
  131. """class analysis for dataset"""
  132. sample_results = simple_analyse(
  133. dataset_path, max_recorded_sample_cnts=float("inf"), show_label=False
  134. )
  135. lab_infos = sample_results[-3] + sample_results[-2] + sample_results[-1]
  136. defaultdict(int)
  137. img_shapes = [] # w, h
  138. ratios_w = []
  139. ratios_h = []
  140. for info in lab_infos:
  141. img = np.asarray(ImageOps.exif_transpose(Image.open(info["img_path"])))
  142. img_h, img_w = np.shape(img)[:2]
  143. img_shapes.append([img_w, img_h])
  144. for box in info["box"]:
  145. box = np.array(box).astype(np.int32).reshape(-1, 2)
  146. box_w, box_h = np.max(box, axis=0) - np.min(box, axis=0)
  147. ratio_w = box_w / img_w
  148. ratio_h = box_h / img_h
  149. ratios_w.append(ratio_w)
  150. ratios_h.append(ratio_h)
  151. m_w_img, m_h_img = np.mean(img_shapes, axis=0) # mean img shape
  152. m_num_box = len(ratios_w) / len(lab_infos) # num box per img
  153. ratio_w = [i * 1000 for i in ratios_w]
  154. ratio_h = [i * 1000 for i in ratios_h]
  155. w_bins = int((max(ratio_w) - min(ratio_w)) // 10)
  156. h_bins = int((max(ratio_h) - min(ratio_h)) // 10)
  157. fig, ax = plt.subplots()
  158. ax.hist(ratio_w, bins=w_bins, rwidth=0.8, color="yellowgreen")
  159. ax.set_xlabel("Width rate *1000")
  160. ax.set_ylabel("number")
  161. canvas = FigureCanvasAgg(fig)
  162. canvas.draw()
  163. width, height = fig.get_size_inches() * fig.get_dpi()
  164. bar_array = np.frombuffer(canvas.tostring_rgb(), dtype="uint8").reshape(
  165. int(height), int(width), 3
  166. )
  167. # pie
  168. fig, ax = plt.subplots()
  169. ax.hist(ratio_h, bins=h_bins, rwidth=0.8, color="pink")
  170. ax.set_xlabel("Height rate *1000")
  171. ax.set_ylabel("number")
  172. canvas = FigureCanvasAgg(fig)
  173. canvas.draw()
  174. width, height = fig.get_size_inches() * fig.get_dpi()
  175. pie_array = np.frombuffer(canvas.tostring_rgb(), dtype="uint8").reshape(
  176. int(height), int(width), 3
  177. )
  178. os.makedirs(output, exist_ok=True)
  179. fig_path = os.path.join(output, "histogram.png")
  180. img_array = np.concatenate((bar_array, pie_array), axis=1)
  181. cv2.imwrite(fig_path, img_array)
  182. return {"histogram": os.path.join("check_dataset", "histogram.png")}
  183. # return {
  184. # "图像平均宽度": m_w_img,
  185. # "图像平均高度": m_h_img,
  186. # "每张图平均文本检测框数量": m_num_box,
  187. # "检测框相对宽度分布图": fig1_path,
  188. # "检测框相对高度分布图": fig2_path
  189. # }