dataset_wtf.py

import os
from collections import defaultdict
import argparse


def wtf(path):
    dirs = [name for name in os.listdir(path) if os.path.isdir(name) and "JESSII" in name or "JESSI" in name]
    dirs.sort()

    images_dict = defaultdict(int)
    labels_dict = defaultdict(int)

    for dirname in dirs:
        for file in os.listdir(f"{path}/{dirname}"):
            if "png" in file:
                images_dict[dirname] += 1
            elif "json" in file:
                labels_dict[dirname] += 1

    total_images = sum(images_dict.values())
    total_labels = sum(labels_dict.values())

    print(f"{'Directory name':<52} | {'Images':6} | {'Labels':6}")
    print("-" * 80)
    for dirname in dirs:
        num_images = images_dict[dirname]
        num_labels = labels_dict[dirname]
        print(f"{dirname:<52} | {num_images:6d} | {num_labels:6d}")
    print("-" * 80)
    print(f"{'Total':<52} | {total_images:6d} | {total_labels:6d}")


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="A small utility for showing the number of images and labels in a directory."
    )
    parser.add_argument("path", nargs="?", help="path to do dataset_wtf on", default=".")
    args = parser.parse_args()

    wtf(args.path)


if __name__ == '__main__':
    main()