YOLOv8 Mask IoU 값 구하기

YOLOv8의 Segment로 예측한 Mask 값과 Labeling한 Mask 값 사이의 IoU 값을 직접 구해봤습니다. ultralytics의 document에는 mask_iou 함수를 사용해서 IoU 값을 구하는 함수가 있었습니다.

위 함수에서는 Tensor 값을 입력하고 Tensor 값을 리턴받는 형태입니다. 해당 함수를 사용하려면 label.txt의 좌표 값을 Tensor로 바꾸고 결과로 나온 (N, M) 형태의 Tensor를 수치 값으로 변환하는 작업이 필요하다고 판단했기 때문에 직접 함수를 작성해서 Mask IoU 값을 구하기로 했습니다.

먼저, 이미지의 예측 mask 값은 아래의 코드로 확인할 수 있습니다.

# model load
model = YOLO('학습한 모델의 경로')
# predict image
results = model(test_image)

results

# 결과
ultralytics.engine.results.Results object with attributes:

boxes: ultralytics.engine.results.Boxes object
keypoints: None
masks: ultralytics.engine.results.Masks object
names: {0: 'no_battery_cover', 1: 'velcro'}
obb: None
orig_img: array([[[ 7,  2,  1],
        [ 7,  2,  1],
        [ 7,  2,  1],
        ...,
        [10,  8,  8],
        [ 9,  7,  7],
        [ 9,  7,  7]],

       [[ 7,  2,  1],
        [ 7,  2,  1],
        [ 7,  2,  1],
        ...,
        [10,  8,  8],
        [ 9,  7,  7],
        [ 9,  7,  7]],

       [[ 7,  2,  1],
        [ 7,  2,  1],
        [ 7,  2,  1],
        ...,
        [10,  8,  8],
        [ 9,  7,  7],
        [ 9,  7,  7]],

       ...,

       [[12, 10, 10],
        [12, 10, 10],
        [12, 10, 10],
        ...,
        [ 5,  0,  1],
        [ 5,  0,  1],
        [ 5,  0,  1]],

       [[13, 11, 11],
        [13, 11, 11],
        [13, 11, 11],
        ...,
        [ 5,  0,  1],
        [ 5,  0,  1],
        [ 5,  0,  1]],

       [[14, 12, 12],
        [14, 12, 12],
        [14, 12, 12],
        ...,
        [ 5,  0,  1],
        [ 5,  0,  1],
        [ 5,  0,  1]]], dtype=uint8)
orig_shape: (640, 640)
path: '/yolov8/datasets/test/images/Screenshot_2024_02_20_16_32_08_jpg.rf.bff8b03c49f87bc3829aae633488f262.jpg'
probs: None
save_dir: 'runs/segment/predict2'
speed: {'preprocess': 5.001306533813477, 'inference': 5.978822708129883, 'postprocess': 3.043651580810547}

# 예측한 Mask data 조회
for i in results:
    results.masks

# 예측 mask의 xy 좌표값
result[0].masks.xyn[0]

예측 mask의 xy 좌표값은 array로 출력 됩니다.

# 출력 결과
array([[    0.41875,       0.425],
       [    0.41719,     0.42656],
       [    0.41563,     0.42656],
       [    0.40781,     0.43437],
       [    0.45625,     0.42656],
       ...
       [    0.45469,     0.42656],
       [    0.45312,       0.425]], dtype=float32)

Roboflow로 labeling한 txt 파일은 아래의 형식으로 작성되어 있습니다.

# label.txt 형식
'1 0.50390625 0.5052083328125 0.494140625 0.4739583328125 0.4755859375 0.44444444375 0.4541015625 0.4270833328125 0.4384765625 0.4236111109375 0.4365234375 0.4270833328125 0.4228515625 0.4270833328125 0.4169921875 0.43055555625 0.41015625 0.4427083328125 0.404296875 0.4704861109375 0.40625 0.4739583328125 0.40625 0.484375 0.412109375 0.5017361109375 0.4345703125 0.5416666671875 0.4462890625 0.5520833328125 0.4658203125 0.559027778125 0.4677734375 0.5625 0.48046875 0.5607638890625 0.4931640625 0.5520833328125 0.501953125 0.5364583328125 0.50390625 0.5260416671875 0.50390625 0.5052083328125'

위의 label.txt 좌표와 예측의 결과로 출력된 mask와 비교해서 알아보겠습니다.

# 예측 mask의 xy 좌표값
result[0].masks.xyn[0]

# Extract x and y coordinates
x = [point[0] for point in result[0].masks.xyn[0]]
y = [point[1] for point in result[0].masks.xyn[0]]

예측 결과의 array를 x, y 리스트를 생성해서 각각 담았습니다.

# Extract x and y coordinates
xx = []
yy = []

# 첫 번째 값은 클래스 정보이기 때문에 index 1에서 부터 시작
# 홀수 번째 있는 값은 x 좌표, 짝수 번째에 있는 값은 y 좌표
for idx, i in enumerate(label_0[1:]):
    if idx % 2 == 0:
        xx.append(float(i))
    else:
        yy.append(float(i))

label.txt의 값을 xx, yy 리스트를 생성해서 각각 담았습니다.

두 mask에 대해서 좌표평면에 점을 찍어 제대호 값이 들어갔는지 확인해보겠습니다.

import matplotlib.pyplot as plt

# Plot the data
plt.figure()
plt.scatter(xx, yy)
plt.scatter(x, y)
plt.xlim(0, 1)
plt.ylim(0, 1)

실제 mask와 예측 mask를 비교해서 값이 알맞게 들어갔는지 확인했으니, 해당 polygon 영역의 IoU 값을 구하겠습니다.

stackoverflow를 참고했습니다.

작성한 함수는 calculate_iou로 input 값으로 model의 결과값 predict와 label.txt의 경로 label을 받습니다.

# predict = model('이미지경로')
# label = 'label.txt경로'
def calculate_iou(predict, label):

    # 변수 선언
    iou_result = 0
    iou_list = []
    label_list = label.split('\n')
    true_mask_cnt = len(label_list)
    predict_mask_cnt = predict[0].boxes.cls.size()[0]

    # 예측 mask와 정답 mask 개수가 같은 경우
    if true_mask_cnt == predict_mask_cnt:

        # 예측 mask와 정답 mask가 없는 경우
        if predict[0].masks==None:
            return 1

        # 같은 수의 mask iou 계산
        for index in range(true_mask_cnt):
            label_mask_x = []
            label_mask_y = []
            predict_mask_x = [point[0] for point in predict[0].masks.xyn[0]]
            predict_mask_y = [point[1] for point in predict[0].masks.xyn[0]]
            label_ele = label_list[index].split()
            for label_idx, p in enumerate(label_ele[1:]):
                if label_idx % 2 == 0:
                    label_mask_x.append(float(p))
                else:
                    label_mask_y.append(float(p))

            predict_mask_point = []
            for pm_idx in range(len(predict_mask_x)):
                predict_mask_point.append((predict_mask_x[pm_idx],predict_mask_y[pm_idx]))

            label_mask_point = []
            for lm_idx in range(len(label_mask_x)):
                label_mask_point.append((label_mask_x[lm_idx],label_mask_y[lm_idx]))

            polygon1 = Polygon(predict_mask_point)
            polygon2 = Polygon(label_mask_point)
            intersect = polygon1.intersection(polygon2).area
            union = polygon1.union(polygon2).area
            iou = intersect / union
            iou_list.append(iou)

        return sum(iou_list)/true_mask_cnt

    # 개수가 다른데 둘 중의 하나가 0인 경우
    if true_mask_cnt==0 or predict_mask_cnt==0:
        return 0


    # 예측 mask와 정답 mask 개수가 다른 경우
    ## 정답 mask 좌표 리스트
    label_mask_point_list = []
    for index in range(true_mask_cnt):
        label_mask_x = []
        label_mask_y = []
        label_ele = label_list[index].split()
        # 정답 mask 개수만큼 loop
        # 정답 mask 개수 == len(label_mask_point_list)

        for label_idx, p in enumerate(label_ele[1:]):
            if label_idx % 2 == 0:
                label_mask_x.append(float(p))
            else:
                label_mask_y.append(float(p))

        label_mask_point = []

        for lm_idx in range(len(label_mask_x)):

            label_mask_point.append((label_mask_x[lm_idx],label_mask_y[lm_idx]))

        label_mask_point_list.append(label_mask_point)

    ## 예측 mask 좌표 리스트
    predict_mask_point_list = []

    for index in range(predict_mask_cnt):
        predict_mask_x = [point[0] for point in predict[0].masks.xyn[index]]
        predict_mask_y = [point[1] for point in predict[0].masks.xyn[index]]
        predict_mask_point = []

        for pm_idx in range(len(predict_mask_x)):
            predict_mask_point.append((predict_mask_x[pm_idx],predict_mask_y[pm_idx]))

        predict_mask_point_list.append(predict_mask_point)

    ## 리스트를 조합으로 IoU값을 구함
    temp_list = product(label_mask_point_list, predict_mask_point_list)
    temp_list_len = 0

    for polygons in temp_list:
        polygon1 = Polygon(polygons[0])
        polygon2 = Polygon(polygons[1])
        intersect = polygon1.intersection(polygon2).area
        union = polygon1.union(polygon2).area
        iou = intersect / union

        iou_result += iou

        temp_list_len += 1

    return iou_result / temp_list_len