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"]

# 1 -> 0, 2 -> 1, 3 -> 2
CHANGE_ID_TABLE = {
    "1" : "0",
    "2" : "1",
    "3" : "2"
}

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 CHANGE_ID_TABLE:
                    line_split[0] = CHANGE_ID_TABLE[line_split[0]]
                
                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))}")