[PyTorch] Object Detection IoU 병렬 구현

설명

IoU

IoU는 Object Detection 및 Segmentation 모델에서 Ground Truth와 Predicted Box가 얼마나 겹치는지 측정하여 정확도를 나타낸다. 위 그림과 같이 두 BBox 사이 교차하는 부분의 비율이 전체 비율의 어느정도 인지를 나타낸다.

$$IoU= \frac{A \cap B}{A \cup B}$$

다른 표현으로는 자카드 지수(Jaccard Index)라고도 하는데, 이는 두 집합 사이의 유사도나 다양성을 측정하는 통계량을 의미한다.

$$J(A,B) = \frac{\vert A \cap B \vert}{\vert A \cup B \vert}$$

코드

코드는 [1]에서 코드를 빌려왔고, Gemini를 활용하여 병렬처리 방식으로 확장했다.

# x1,y1,x2,y2
box_a = [3, 3, 6, 6]
box_b = [4, 4, 7, 7]

2차원 좌표평면 위에 최소 (3, 3)에서 최대 (6, 6)을 가지는 box_a와 (4, 4)에서 (7, 7)을 가지는 box_b 있다고 하자.

이를 matplot을 통해 시각화하면 다음과 같이 그려진다.

box_a와 box_b

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

plt.scatter(box_a[::2], box_a[1::2])
plt.scatter(box_b[::2], box_b[1::2])

plt.gca().add_patch(
    patches.Rectangle(
        (box_a[:2]), box_a[2] - box_a[0], box_a[3] - box_a[1], color="tab:blue", alpha=0.5
    )
)
plt.gca().add_patch(
    patches.Rectangle(
        (box_b[:2]), box_b[2] - box_b[0], box_b[3] - box_b[1], color="tab:orange", alpha=0.5
    )
)

plt.grid()
plt.xticks(np.arange(0, 10, 1))
plt.yticks(np.arange(0, 10, 1))
plt.tight_layout()
plt.show()

box의 각 영역의 넓이를 구하면 다음과 같다.

box_a_area = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
box_b_area = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
print(box_a_area, box_b_area) # 9 9

해당 식은 (x_2 - x_1) * (y_2 * y_1)으로 (가로 길이) * (세로 길이)를 나타낸다.

참고로 [1] 자료에서는 각 지름에 (box_a[2] - box_a[0]) 대신 (box_a[2] - box_a[0] + 1)을 사용하였는데,
이는,
1. 넓이가 0이 되는 것을 방지 하기 위해
2. 픽셀을 점이 아닌 네모칸으로 보고 있어서 (적어도 1 길이를 가지는)
로 알려져 있다. (Gemini 피셜)

그래서 넓이 관점에서 정확하지 않기 때문에, +1은 제외하여 계산했다.

inter_x1 = max(box_a[0], box_b[0])
inter_y1 = max(box_a[1], box_b[1])
inter_x2 = min(box_a[2], box_b[2])
inter_y2 = min(box_a[3], box_b[3])

inter_w = max(0, inter_x2 - inter_x1)
inter_h = max(0, inter_y2 - inter_y1)

inter_area = inter_w * inter_h
print(inter_area) # 4

교차되는 사각형의 왼쪽 최소점(x1, 왼쪽 아래)은 왼쪽 아래에서 시작점이 되기 때문에 두 x1 중 큰 점이 된다. (max)
교차되는 사각형의 왼쪽 최대점(y1, 왼쪽 위)은 왼쪽 위에서 시작점이 되기 때문에 두 y1 중 큰 점이 된다. (max)
교차되는 사각형의 오른쪽 최소점(x2, 오른쪽 아래)은 오른쪽 아래에서 끝점이 되기 때문에 두 x2 중 작은 점이 된다. (min)
교차되는 사각형의 오른쪽 최대점(y2, 오른쪽 위)은 오른쪽 위에서 끝점이 되기 때문에 두 y2 중 작은 점이 된다. (min)

교차되는 사각형의 가로는 x2-x1의 절대값이 되고,
교차되는 사각형의 세로는 y2-y1의 절대값이 된다.

공식에 의해 IoU 값은 다음과 같이 나온다.

iou = inter_area / (box_a_area + box_b_area - inter_area) # 4 / (9 + 9 - 4)
print(iou) # 0.2857142857142857

---

Object Detection 시 BBox와 GT의 IoU 계산은 필수적이다. src_box(BBox)와 dst_box(GT)가 각각 여러 개 존재 할 때, 이를 하나씩 게산하게 되면 연산의 오버헤드가 발생하여 속도가 느려질 것이다. 예를들어, 2개의 bbox가 3개의 GT와 score 계산을 위해 모든 IoU가 필요할 수 있다.

이를 위해 torch를 사용하여 2개, 3개 임의의 box를 만들어본다.

import torch

box_a = torch.tensor(
    [
        [0, 0, 5, 5],
        [8, 8, 10, 10],
    ],
    dtype=torch.float32,
)

box_b = torch.tensor(
    [
        [0, 2, 1, 6],
        [2, 2, 6, 4],
        [9, 9, 11, 11],
    ],
    dtype=torch.float32,
)

이를 matplot을 통해 시각화 하면 다음과 같다.

boxes_a와 boxes_b

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

plt.scatter(box_a[:,::2], box_a[:,1::2])
plt.scatter(box_b[:,::2], box_b[:,1::2])

for box in box_a:
    plt.gca().add_patch(
        patches.Rectangle(
            (box[:2]), box[2] - box[0], box[3] - box[1], color="tab:blue", alpha=0.5
        )
    )

for box in box_b:
    plt.gca().add_patch(
        patches.Rectangle(
            (box[:2]), box[2] - box[0], box[3] - box[1], color="tab:orange", alpha=0.5
        )
    )

plt.grid()
plt.xticks(np.arange(0, 12, 1))
plt.yticks(np.arange(0, 12, 1))
plt.tight_layout()
plt.show()

참고로 patches 라이브러리로 도형을 그릴 때 병렬로 안되서 for문으로 각각 그려 주었다.

box_a_area = (box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])
box_b_area = (box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])

print(box_a_area, box_b_area) # tensor([25.,  4.]) tensor([4., 8., 4.])

python의 list slicing 문법으로 처음에 한 쌍으로 했을 때 코드에서,
행 부분에 `:`을 붙여주면 된다. 이러면 전체행의 n번째 열만 가져오라는 뜻이 된다.

각 쌍에 대해서 (가로 길이) * (세로 길이) 넓이 연산을 해준 후 저장한다.

이제 각 쌍에 대해서 교차 영역을 계산해 주어야 한다. 교차되지 않는 것도 포함해서 2x3 개 나와야 한다. 

교차되는 사각형의 왼쪽 최소점(x1, 왼쪽 아래)은 왼쪽 아래에서 시작점이 되기 때문에 두 x1 중 큰 점이 된다. (max)
교차되는 사각형의 왼쪽 최대점(y1, 왼쪽 위)은 왼쪽 위에서 시작점이 되기 때문에 두 y1 중 큰 점이 된다. (max)
교차되는 사각형의 오른쪽 최소점(x2, 오른쪽 아래)은 오른쪽 아래에서 끝점이 되기 때문에 두 x2 중 작은 점이 된다. (min)
교차되는 사각형의 오른쪽 최대점(y2, 오른쪽 위)은 오른쪽 위에서 끝점이 되기 때문에 두 y2 중 작은 점이 된다. (min)

교차되는 사각형의 가로는 x2-x1의 절대값이 되고,
교차되는 사각형의 세로는 y2-y1의 절대값이 된다.

이 규칙을 적용해주어야 하기 때문에 max, min을 구해주어야 하는데, 현재 box_a의 차원은 (2, 4)와 box_b의 차원은 (3, 4)이다.

이 때 python의 broadcasting 방법으로 차원을 확장하여 비교할 수 있다.

(2, 4) - (3 ,4) 이니까, (2, 1, 4) - (1, 3, 4)처럼 차원을 늘려준다.
2번째(0-index) 축을 고정해놓고 이를 생략하고 보았을 때, (2, 1) - (1, 3)이 된다.

행렬 곱 (m x p) * (p x n) = (m x n)처럼 앞 행렬의 열 개수와 뒷 행렬의 행 개수가 동일해야 행렬곱이 가능해진다.
브로드캐스팅으로 행렬곱 뿐만 아니라, 각 요소에 대해 비교 연산이 가능해진다. 이는 내부적으로 이중 반복문을 병렬로 수행하도록 되어 있다.

따라서 `unsqueeze` 메소드를 통해 차원을 각 축에 확장하면 된다.
unsqueeze(0)은 0번째에 차원을 추가하고, unsqueeze(1)은 1번째 차원을 추가한다.

print(box_a.unsqueeze(1).shape) # torch.Size([2, 1, 4])
print(box_b.unsqueeze(0).shape) # torch.Size([1, 3, 4])

inter_x1 = torch.max(box_a[:, 0].unsqueeze(1), box_b[:, 0].unsqueeze(0))
print(inter_x1) # tensor([[0., 2., 9.],
                # [8., 8., 9.]])

inter_y1 = torch.max(box_a[:, 1].unsqueeze(1), box_b[:, 1].unsqueeze(0))
print(inter_y1) # tensor([[2., 2., 9.],
                #[8., 8., 9.]])

inter_x2 = torch.min(box_a[:, 2].unsqueeze(1), box_b[:, 2].unsqueeze(0))
print(inter_x2) # tensor([[ 1.,  5.,  5.],
                # [ 1.,  6., 10.]])

inter_y2 = torch.min(box_a[:, 3].unsqueeze(1), box_b[:, 3].unsqueeze(0))
print(inter_y2) # tensor([[ 5.,  4.,  5.],
                # [ 6.,  4., 10.]])

그림과 하나씩 천천히 비교해보면 교차하는 영역의 좌표가 구해졌음을 확인할 수 있다.

inter_w = (inter_x2 - inter_x1).clamp(min=0)
inter_h = (inter_y2 - inter_y1).clamp(min=0)
inter_area = inter_w * inter_h

print(inter_w) # tensor([[1., 3., 0.], [0., 0., 1.]])
print(inter_h) # tensor([[3., 2., 0.], [0., 0., 1.]])
print(inter_area) # tensor([[3., 6., 0.], [0., 0., 1.]])

clamp 메소드를 통해 0이하가 되는 텐서는 0이 되도록 맞춰 주었다.

print 결과를 보면 교차되는 영역이 총 3개임을 알 수 있다.
A의 1번째 box는 B의 1, 2번째 박스와 교차하고 A의 2번째 박스는 3번째 박스와 교차한다는 사실을 알고 있다.
나머지 조합은 교차하지 않아 0이 되었다.

ious = inter_area / (box_a_area.unsqueeze(1) + box_b_area.unsqueeze(0) - inter_area)
ious # tensor([[0.1154, 0.2222, 0.0000], [0.0000, 0.0000, 0.1429]])

공식을 통해 병렬로 각 IoU를 구해주었다. box_a_area와 box_b_area는 각 조합을 구할 때 차원이 맞지 않아, 브로드캐스팅을 해준 후 경우의 수 2x3개가 나오도록 해야한다.

참고자료

1. minimin2. (2021.02.08). [딥러닝] IoU 설명, python 코드(Intersection over Union, object detection 평가방법). 기록은 기억을 지배한다. https://minimin2.tistory.com/144

2. Gemini