YOLOv5 이미지 데이터셋을 구축해보자! - labelImg 사용법

컴퓨터 비전 관련 프로젝트를 하다보면, 데이터셋을 직접 구축해야할 일이 꽤 생깁니다. labelImg라는 사용하기 쉬운 라이브러리를 사용해서 직접 Bounding Box를 그려 데이터셋을 생성해봅시다.

 

① labelImg 깃허브에서 다운로드 받기

https://github.com/tzutalin/labelImg

GitHub - tzutalin/labelImg: 🖍️ LabelImg is a graphical image annotation tool and label object bounding boxes in images

먼저 위 깃허브 링크로 들어가시면, 아래와 같은 화면을 보실 수 있습니다!

 

초록색 Code를 클릭하신 후 링크를 복사해서 clone 하셔도 되고, 깃허브에 익숙하지 않으신 분들은 ZIP파일을 직접 다운로드 후 압축 풀어주셔도 무방합니다.

 

 

② predefined_classes.txt 수정하기

data 폴더로 들어가시면 predefined_classes라는 txt파일을 보실 수 있습니다. 편하게 Bounding Box를 그리기 위해선 이 파일을 수정해야 합니다. txt파일에 클래스명을 입력해주시면 되는데요, 예를 들어 Apple, Banana, Tomato 클래스가 있다면 아래 이미지 처럼 수정해주시면 되겠습니다.

 

 

③ labelImg.py 실행하기 

저는 Powershell prompt를 통해 실행하였습니다. ls와 cd 명령어를 사용하여 labelImg.py가 있는 폴더로 이동합니다. 이후 아래와 같은 명령어를 실행하여 실행해주면 됩니다!

python labelImg.py

 

 

④ 설정

정상적으로 실행하셨다면 위와 같은 화면을 보실 수 있으실 겁니다. 이제 Bounding Box를 그리기 위한 설정을 해보겠습니다.

 

(1) Save directory 설정

클래스와 Bounding Box 좌표를 담고있는 txt파일이 저장될 폴더를 설정해주셔야 합니다. 왼쪽에 Change Save Dir라는 버튼을 보실 수 있는데요, 간단히 클릭 후에 저장 공간을 설정해주시면 됩니다!

 

(2) YOLO 로 설정

어느 딥러닝 모델을 사용하느냐에 따라서 Bounding Box 등이 어떤 형태의 파일로 저장될 지가 달라집니다.  저 같은 경우 YOLO 모델을 사용할 것이기 때문에, 왼쪽 Save버튼 아래에 있는 버튼을 클릭하면서 YOLO가 나올때까지 클릭해주시면 됩니다. (위 이미지 처럼 YOLO라는 글자의 버튼이 화면에 보이셔야 합니다. YOLO로 변경하겠다고 한번 더 클릭하시면 안됩니다.)

 

(3) Open Directory

이제 Bounding Box를 설정할 이미지 데이터가 있는 폴더를 눌러주시면 됩니다. 왼쪽 Open Dir를 클릭해서, 이미지 데이터가 있는 폴더를 열어주시면 됩니다. 하나의 이미지만 Bounding Box를 그리실거면 Open을 클릭하셔도 무방합니다. 

 

 

⑤ Bounding Box 그리기

이제 다 됐네요. 모든 준비는 끝났으니 마우스를 이용해 바운딩 박스를  직접 그리면 됩니다! 

왼쪽에 있는 버튼 중 Create RectBox를 클릭해주셔도 되고, w라는 단축키로 RectBox를 생성해주셔도 무방합니다. (단축키 사용이 훨씬 편하겠죠?) 

이제 마우스 드래그를 하여 Bounding Box를 그리면 아래 이미지와 같은 상황일겁니다.

 

사과 이미지 출처 : https://www.kaggle.com/datasets/kritikseth/fruit-and-vegetable-image-recognition

오른쪽 창에 나와있는 클래스들은, 저희가 앞서 predefined_classes.txt에서 입력한 클래스들입니다. 클래스 중 하나를 클릭하신 후 OK를 클릭해주면 됩니다! 만약 하나의 이미지 데이터 내에 여러 객체가 존재한다면, 똑같이 Bounding Box를 그려주신 후, 해당하는 클래스를 선택하고 OK를 클릭해주시면 됩니다. OK를 클릭하시고 나면 아래와 같은 화면일겁니다.

 

이제 다 끝났네요. 왼쪽에 있는 Save 버튼을 눌러주시면, 앞서 설정한 Save Directory에 txt파일이 생성되고, txt파일의 이름은 사용한 이미지의 이름과 같습니다. (YOLO의 경우 txt파일이고, 다른 딥러닝 모델로 설정해주셨다면 txt파일이 아닐 수도 있습니다.) 

 

 

생성된 txt파일을 열어보니 위와 같네요. 5개의 숫자가 띄어쓰기로 구분되어 있는 것을 보실 수 있습니다. 맨 앞 숫자는 Class number, 그리고 나머지 4개의 숫자는 각각 x_start, y_start, width, height입니다. 그리고 Bounding Box의 좌표는 0~1 사이의 값으로, 이미지 전체에 대한 상대 좌표(?)라는 것에 주의하면 좋을 것 같습니다!

 

 

 

 

잘못된 내용에 대한 피드백은 언제나 환영합니다.

 

 

[다음글]

https://aiclaudev.tistory.com/30

YOLOv5 이미지 데이터셋을 구축해보자! - Imgaug 사용법 (Bounding Box 같이 회전시키기)