| from glob import glob | |
| from typing import Tuple,List | |
| import os | |
| import argparse | |
| import json | |
| from matplotlib import pyplot as plt | |
| import numpy as np | |
| from PIL import Image | |
| def get_args()->Tuple[str, str]: | |
| """ | |
| Return: | |
| --dataset_dir: dataset dir. | |
| --save_dir: save dir. | |
| """ | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--dataset_dir', type=str, default='data/grass') | |
| parser.add_argument('--save_dir', type=str, default='dataset_num_analysis.png') | |
| args = parser.parse_args() | |
| return args.dataset_dir, args.save_dir | |
| def get_mask_files(dataset_dir: str)->List[str]: | |
| """ | |
| get mask files from dataset dir. | |
| Args: | |
| dataset_dir: dataset dir. | |
| Return: | |
| mask_filenames: list of mask filenames. | |
| """ | |
| mask_filenames = glob(os.path.join(dataset_dir, "ann_dir", "*", "*.png")) | |
| return mask_filenames | |
| def main(): | |
| dataset_dir, save_dir = get_args() | |
| mask_filenames = get_mask_files(dataset_dir) | |
| statistic = {} | |
| for mask_filename in mask_filenames: | |
| mask = np.array(Image.open(mask_filename)) | |
| classes = np.unique(mask) | |
| for class_ in classes: | |
| class_ = int(class_) | |
| if class_ not in statistic: | |
| statistic[class_] = 0 | |
| statistic[(class_)] += int(np.sum(mask == class_)) | |
| classes = list(statistic.keys()) | |
| clasees_num = list(statistic.values()) | |
| plt.title("Dataset Analysis") | |
| bars = plt.bar(classes, clasees_num) | |
| for bar in bars: | |
| height = bar.get_height() | |
| plt.text(bar.get_x() + bar.get_width() / 2, height + 5, str(height), ha='center', va='bottom') | |
| plt.savefig(save_dir,dpi=300) | |
| if __name__ == "__main__": | |
| main() |