analyse_dataset.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  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 pycocotools.coco import COCO
  21. from .....utils.fonts import PINGFANG_FONT_FILE_PATH
  22. def deep_analyse(dataset_dir, output):
  23. """class analysis for dataset"""
  24. tags = ["train", "val"]
  25. all_instances = 0
  26. for tag in tags:
  27. annotations_path = os.path.abspath(
  28. os.path.join(dataset_dir, f"annotations/instance_{tag}.json")
  29. )
  30. labels_cnt = defaultdict(list)
  31. coco = COCO(annotations_path)
  32. cat_ids = coco.getCatIds()
  33. for cat_id in cat_ids:
  34. cat_name = coco.loadCats(ids=cat_id)[0]["name"]
  35. labels_cnt[cat_name] = labels_cnt[cat_name] + coco.getAnnIds(catIds=cat_id)
  36. all_instances += len(labels_cnt[cat_name])
  37. if tag == "train":
  38. cnts_train = [len(cat_ids) for cat_name, cat_ids in labels_cnt.items()]
  39. elif tag == "val":
  40. cnts_val = [len(cat_ids) for cat_name, cat_ids in labels_cnt.items()]
  41. classes = [cat_name for cat_name, cat_ids in labels_cnt.items()]
  42. sorted_id = sorted(
  43. range(len(cnts_train)), key=lambda k: cnts_train[k], reverse=True
  44. )
  45. cnts_train_sorted = sorted(cnts_train, reverse=True)
  46. cnts_val_sorted = [cnts_val[index] for index in sorted_id]
  47. classes_sorted = [classes[index] for index in sorted_id]
  48. x = np.arange(len(classes))
  49. width = 0.5
  50. # bar
  51. os_system = platform.system().lower()
  52. if os_system == "windows":
  53. plt.rcParams["font.sans-serif"] = "FangSong"
  54. else:
  55. font = font_manager.FontProperties(fname=PINGFANG_FONT_FILE_PATH)
  56. fig, ax = plt.subplots(figsize=(max(8, int(len(classes) / 5)), 5), dpi=120)
  57. ax.bar(x, cnts_train_sorted, width=0.5, label="train")
  58. ax.bar(x + width, cnts_val_sorted, width=0.5, label="val")
  59. plt.xticks(
  60. x + width / 2,
  61. classes_sorted,
  62. rotation=90,
  63. fontproperties=None if os_system == "windows" else font,
  64. )
  65. ax.set_ylabel("Counts")
  66. plt.legend()
  67. fig.tight_layout()
  68. fig_path = os.path.join(output, "histogram.png")
  69. fig.savefig(fig_path)
  70. return {"histogram": os.path.join("check_dataset", "histogram.png")}