기계는 거짓말하지 않는다

YOLO label 텍스트 classid 변경 본문

AI

YOLO label 텍스트 classid 변경

KillinTime 2022. 10. 14. 13:30

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))}")

 

'AI' 카테고리의 다른 글

YOLO label 텍스트 object 개수  (2) 2022.10.23
YOLO label 텍스트 classid 제외  (1) 2022.10.23
YOLO label to Labelme JSON label 변환  (0) 2022.10.02
YOLOv3 deepSORT KeyError: "The name 'net/images:0'  (0) 2022.07.01
min, max 좌표 YOLO label format 변환  (0) 2022.05.27
'AI' Related Articles more
Comments