analyse_dataset.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  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 os.path as osp
  13. import matplotlib.pyplot as plt
  14. import numpy as np
  15. from PIL import Image, ImageOps
  16. from .....utils.file_interface import custom_open
  17. from .....utils.logging import info
  18. def anaylse_dataset(dataset_dir, output):
  19. """class analysis for dataset"""
  20. split_tags = ["train", "val"]
  21. label2count = {tag: dict() for tag in split_tags}
  22. for tag in split_tags:
  23. mapping_file = osp.join(dataset_dir, f"{tag}.txt")
  24. with custom_open(mapping_file, "r") as fp:
  25. lines = filter(None, (line.strip() for line in fp.readlines()))
  26. for i, line in enumerate(lines):
  27. _, ann_file = line.split(" ")
  28. ann_file = osp.join(dataset_dir, ann_file)
  29. ann = np.array(
  30. ImageOps.exif_transpose(Image.open(ann_file)), "uint8")
  31. for idx in set(ann.reshape([-1]).tolist()):
  32. if idx == 255:
  33. continue
  34. if idx not in label2count[tag]:
  35. label2count[tag][idx] = 1
  36. else:
  37. label2count[tag][idx] += 1
  38. if label2count[tag].get(0, None) is None:
  39. label2count[tag][0] = 0
  40. train_label_idx = np.array(list(label2count["train"].keys()))
  41. val_label_idx = np.array(list(label2count["val"].keys()))
  42. label_idx = np.array(list(set(train_label_idx) | set(val_label_idx)))
  43. x = np.arange(len(label_idx))
  44. train_list = []
  45. val_list = []
  46. for i in range(len(label_idx)):
  47. train_list.append(label2count["train"].get(i, 0))
  48. val_list.append(label2count["val"].get(i, 0))
  49. fig, ax = plt.subplots(
  50. figsize=(max(8, int(len(label_idx) / 5)), 5), dpi=120)
  51. width = 0.5,
  52. ax.bar(x, train_list, width=width, label="train")
  53. ax.bar(x + width, val_list, width=width, label="val")
  54. plt.xticks(x + 0.25, label_idx)
  55. ax.set_xlabel('Label Index')
  56. ax.set_ylabel('Sample Counts')
  57. plt.legend()
  58. fig.tight_layout()
  59. fig_path = os.path.join(output, "histogram.png")
  60. fig.savefig(fig_path)
  61. return {"histogram": os.path.join("check_dataset", "histogram.png")}