YOLO label 텍스트 object 개수

classid, x center, y center, width, height 형식의

YOLO label 텍스트 파일의 class 별 object 개수를 확인하는 경우

import os
import glob
import datetime
import time
import copy

# YOLO dataset 구조 형식일 경우
dir_root_path = "D:/Datasets/yolo_split_dataset/labels/"
sub_dir_name_list = ["train", "val"]

name_dict = {}
read_count = 0

start_time = time.time()

for sub_dir_name in sub_dir_name_list:
    label_path_list = glob.glob(os.path.join(dir_root_path + sub_dir_name, "*.txt"))
    read_count = 0
    class_count_dict = {}

    print(f"{sub_dir_name} label")
    
    for label_path in label_path_list:
        with open(label_path, "r") as f:
            lines = f.readlines()
            
            for line in lines:
                line_split = line.split(" ")
                if line_split[0] in class_count_dict:
                    class_count_dict[line_split[0]] += 1
                else:
                    class_count_dict[line_split[0]] = 1
        read_count += 1
        
        if read_count % 1000 == 0:
            print(f"{read_count} / {len(label_path_list)}")
            
    name_dict[sub_dir_name] = copy.deepcopy(class_count_dict)

print('\n< Result >')

# sub directory의 class 별 object 개수
for name in name_dict.keys():
    print(f"[{name}]")
    key_list = list(name_dict[name].keys())
    key_list.sort(key=lambda x: int(x))

    print(key_list)
    for key in key_list:
        print("class:", key, "count:", name_dict[name][key])
    print('')
    
print(f"Total Time: {datetime.timedelta(seconds=int(time.time() - start_time))}")