analyse_dataset.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import platform
  13. from collections import defaultdict
  14. from pathlib import Path
  15. import matplotlib.pyplot as plt
  16. import numpy as np
  17. from matplotlib import font_manager
  18. from pycocotools.coco import COCO
  19. from .....utils.fonts import PINGFANG_FONT_FILE_PATH
  20. def deep_analyse(dataset_dir, output):
  21. """class analysis for dataset"""
  22. tags = ['train', 'val']
  23. all_instances = 0
  24. for tag in tags:
  25. annotations_path = os.path.abspath(
  26. os.path.join(dataset_dir, f'annotations/instance_{tag}.json'))
  27. labels_cnt = defaultdict(list)
  28. coco = COCO(annotations_path)
  29. cat_ids = coco.getCatIds()
  30. for cat_id in cat_ids:
  31. cat_name = coco.loadCats(ids=cat_id)[0]["name"]
  32. labels_cnt[cat_name] = labels_cnt[cat_name] + coco.getAnnIds(
  33. catIds=cat_id)
  34. all_instances += len(labels_cnt[cat_name])
  35. if tag == 'train':
  36. cnts_train = [
  37. len(cat_ids) for cat_name, cat_ids in labels_cnt.items()
  38. ]
  39. elif tag == 'val':
  40. cnts_val = [
  41. len(cat_ids) for cat_name, cat_ids in labels_cnt.items()
  42. ]
  43. classes = [cat_name for cat_name, cat_ids in labels_cnt.items()]
  44. sorted_id = sorted(
  45. range(len(cnts_train)), key=lambda k: cnts_train[k], reverse=True)
  46. cnts_train_sorted = sorted(cnts_train, reverse=True)
  47. cnts_val_sorted = [cnts_val[index] for index in sorted_id]
  48. classes_sorted = [classes[index] for index in sorted_id]
  49. x = np.arange(len(classes))
  50. width = 0.5
  51. # bar
  52. os_system = platform.system().lower()
  53. if os_system == "windows":
  54. plt.rcParams['font.sans-serif'] = 'FangSong'
  55. else:
  56. font = font_manager.FontProperties(fname=PINGFANG_FONT_FILE_PATH)
  57. fig, ax = plt.subplots(figsize=(max(8, int(len(classes) / 5)), 5), dpi=120)
  58. ax.bar(x, cnts_train_sorted, width=0.5, label='train')
  59. ax.bar(x + width, cnts_val_sorted, width=0.5, label='val')
  60. plt.xticks(
  61. x + width / 2,
  62. classes_sorted,
  63. rotation=90,
  64. fontproperties=None if os_system == "windows" else font)
  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")}