python으로 yolov5 label을 교체하는 script 만들기

오늘은 python을 이용해서 label을 교체하는 script를 만들어보도록 하겠습니다.

yolov5의 경우, label 값은 Integer값으로 들어오게 됩니다. (0, 1, 등.. )

실제로, 0으로 label되어야 하는 경우 1로 label 되는 경우, 어떻게 이를 바꿀 수 있을지 

python script를 짜 보도록 하겠습니다.

 

# 예를 들어, 이렇게 label 되어 있는 파일을

1 0.188281 0.603704 0.050521 0.205556

이렇게 바꿔주고 싶은 경우, 그리고 안에 있는 모든 파일들을 찾아 (txt 파일의 경우) 

바꿔주는 script를 짜 보도록 하겠습니다.

0 0.188281 0.603704 0.050521 0.205556

change_label.py

 

import os

if __name__ == "__main__":
    root_dir = "./target_dir/"
    for (root, dirs, files) in os.walk(root_dir):
        print("# root : " + root)
        if len(dirs) > 0:
            for dir_name in dirs:
                pass
                #print("dir: " + dir_name)

        if len(files) > 0:
            for file_name in files:
                full_name = root + "/" + file_name
                #with open(full_name, "r+") as f:
                #    data = f.read()
                #    f.seek(0)
                    

                    
                #print("dir : " + root)
                if "txt" in file_name:
                    with open(full_name, "r+") as f:
                        lines = f.readlines()
                        new_lines = []
                        for i in range(len(lines)):
                            print(f"line {i} : ", lines[i])
                            if len(lines[i]) > 0 and lines[i][0] == "1":
                                this_line = "0"+lines[i][1:]
                                new_lines.append(this_line)
                        f.seek(0)
                        f.truncate()        
                        f.write("/n".join(new_lines))
                        
                    print(full_name)
                    #print("txt : " + root + "/"+ file_name)
                #print("file: " + file_name)

 

os.walk 를 통해서 file들과 directory들을 iterate한 이후,

f.seek, f.truncate를 통해 있는 내용들을 지우고 

f.write을 통해 써주는 방법을 사용했습니다.