IT Dictionary

YOLO v4 바운딩 박스 좌표 저장하기 본문

Lab

YOLO v4 바운딩 박스 좌표 저장하기

LeU 2021. 6. 28. 22:20

https://keyog.tistory.com/22

 

YOLO v4 custom데이터 훈련하기

이전에 YOLO v4를 설치해서 example실행 까지 잘 마쳤다면, 이제 본인의 데이터를 훈련시키고 응용하는 것까지 알아보도록 하자. YOLO v4 설치에 관해서는 이전글 참조. 2020/05/19 - [Computer Vision/Object det.

keyog.tistory.com

 

지금까지 yolov5를 이용하였지만 deepsort를 사용하기 위해 yolov4를 사용하기로 하였습니다.

간단하게 yolov4-deepsort를 이용하여 검출한 예시입니다.

위 링크등을 토대로 학습을 하고 검출을 하시면 되는데 저는 검출과정에서 바운딩박스의 좌표가 필요합니다.

 

따라서, 이번에는 좌표를 txt파일에 저장하는 과정입니다. 

간단하기 때문에 먼저 deepsort부터 살펴보면

https://github.com/theAIGuysCode/yolov4-deepsort

 

theAIGuysCode/yolov4-deepsort

Object tracking implemented with YOLOv4, DeepSort, and TensorFlow. - theAIGuysCode/yolov4-deepsort

github.com

 

 

yolov4-deepsort 중에서 검출할때 실행되는 object tracker.py를 봐야합니다.

 

해당 부분에서 frame 을 구분하고 있음을 확인하였고,

 

위 부분에서 yolo를 통해 검출된 바운딩박스의 좌표가 bbox에 저장이되고 bbox[0] 이런식으로 접근하고 있음을 확인할 수 있습니다.

 

이를 그대로 활용하여 저는 tracker_log.txt를 생성해주겠습니다.

105라인에 다음을 추가하여 프레임을 구분해주고

 with open('Tracker_log.txt', 'a') as f:
            f.write('Frame : {}'.format(frame_num)+"\n")

 

222라인에 다음을 추가하면 좌표를 저장할 수 있습니다.

# Write txt file
                with open('Tracker_log.txt', 'a') as f:
                    f.write("Tracker ID: {}, {}".format(str(track.track_id), (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])))+"\n")

 

이제 저장하고 검출해보면 자동으로 txt파일이 생성됩니다. 문자열을 자신에 맞게 변경하여 사용할 수 있습니다.