YOLO label 텍스트 classid 제외

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

YOLO label 텍스트 파일의 특정 classid를 제외해야 하는 경우

import os
import glob
import datetime
import time

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

# 제외할 classid
skip_classid_list = ["2", "5", "9"]

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"))
    count = 0

    print(f"{sub_dir_name} label")
    for label_path in label_path_list:
        new_contents = ""
        with open(label_path, "r") as f:
            lines = f.readlines()
            
            for line in lines:
                line_split = line.split(" ")
                if line_split[0] in skip_classid_list:
                    continue
                
                new_contents += (" ".join(line_split))
                
        with open(label_path, "w") as f:
            f.write(new_contents)
        count += 1
    print("complete", count)
    
print(f"Total Time: {datetime.timedelta(seconds=int(time.time() - start_time))}")