Yolo 좌표형식으로 Opencv에서 사각형 그리기

#1 Yolo 좌표 포맷 알아보기

 

제가 원한 데이터셋은

Raw data(깨끗한 원본 사진)  -  GroundTruth data(박스로 레이블이 된 사진) 

이런 셋이었지만...

항공촬영 데이터가 생각보다 너무 적어서 아래와 같은 데이터셋을 받았습니다.

이미지와 텍스트 문서가 같은 이름 쌍으로 이루어진 데이터셋입니다.

102장짜리 데이터쌍

 

텍스트 문서에는 이미지에 있는 객체의 박스좌표가 입력되어 있습니다.

한 줄에 5가지 정보가 적혀있는데 이건 Yolo에서 쓰는 좌표 포맷입니다.

띄어쓰기를 구분자로 해서 <object-class> <x_center> <y_center> <width> <height> 입니다.

<x_center>와 <y_center>는 직사각형의 중심을 표현합니다.

<width>와 <height>는 중심에서부터 사각형의 폭과 높이를 상대적으로 표현합니다.

 

0~1 사이의 값으로 정규화된 수치라서 좌표 포맷이 다르므로

전처리 없이 Opencv의 rectangle 함수에 넣을 수 없습니다.

 

 

OpenCV에서 rectangle 함수로 사각형을 그릴 때는 다음 형식에 맞춰야 합니다.

cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, thickness)

img는 박스가 그려질 원본 이미지 파일입니다.

xmin, ymin 은 그리고 싶은 직사각형의 왼쪽 상단 꼭짓점입니다.

xmax, ymax 은 오른쪽 하단 꼭짓점입니다.

※ 이건 좌표 시작점에 따라 다른데 xmin, ymin에 시작꼭짓점 xmax, ymax는 반대편 꼭짓점을 입력해주시면 됩니다.

color는 그려질 직사각형의 색깔입니다. BGR 순서로 써주세요.

예를 들어 (0,255,0)이라고 하면 초록색 박스가 그려집니다.

thickness는 박스를 그리는 선의 두께입니다.

숫자가 클수록 두꺼운 선으로 박스를 그립니다.

 

#2 Ground Truth 이미지 만들기

자세한 주석을 달아둔 파이썬 코드로 내용을 대체하겠습니다.

import cv2
import numpy as np
import os

#이미지와 레이블 파일들의 경로를 각각 불러와서 정렬합니다.
img_path = os.listdir('./Images/') 
txt_path = os.listdir('./Labels/')
img_path.sort()
txt_path.sort()

for i in img_path:
     img_path = './Images/' + i
     #.jpg를 .txt로 바꿔서 읽는 편법...
     i = i.replace("jpg","txt")
     txt_path = './Labels/' + i

     #Yolo 좌표가 적혀있는 텍스트 파일을 읽어옵니다.
     txt_file = open(txt_path,'r')  
     while True:
          image = cv2.imread(img_path, cv2.IMREAD_COLOR)
          dh,dw,_ = image.shape
          yolo_box = txt_file.readline()

          #텍스트 파일에서 더 읽어올 게 없다면 while 구문을 나갑니다.
          if not yolo_box:  
               break

          #텍스트 중에서 첫번째 인자만 빼고 나머지를 취합니다.
          _, x, y, w, h = map(float, yolo_box.split()) 

          #yolo 좌표가 해상도에 대한 비율이라 뒤에 해상도를 곱합니다.
          xmin = int((x-w / 2) * dw) 
          xmax = int((x+w / 2) * dw)
          ymax = int((y-h / 2) * dh)
          ymin = int((y+h / 2) * dh)

          # 예외처리...
          if xmin < 0 :
               xmin = 0
          if xmax > dw-1 :
               xmax = dw-1
          if ymax < 0 :
               ymax = 0
          if ymin > dh-1 :
              ymin = dh-1

          # 이미지에 초록색 직사각형을 그립니다.
          cv2.rectangle(image,(xmin,ymin),(xmax,ymax),(0,255,0),1)
          
          #저장하고 while문을 돌면서 같은 파일에 이어그립니다.
          cv2.imwrite(img_path,image)

     txt_file.close() #텍스트 파일이 끝나면 닫습니다.

결과물

102장을 1분 조금 넘게 걸려서 다 그려냈습니다.

근데 레이블이 좀 마음에 들지 않는다.

 

이제 제 모델이 추론을 마치고 내뱉은 이미지와

위에서 만든 Ground Truth 이미지를

Matlab으로 비교해서 IoU를 측정코자 합니다.