如何用Python计算YOLOv8多边形分割图像的IoU？

问题描述 #

我需要计算YOLOv8分割模块生成的多边形分割结果的IoU，现有YOLOv8输出的预测掩码，以及通过LabelMe标注JSON转换的真值掩码（用shape_to_mask()转换）。尝试了以下代码，但所有图像的IoU计算结果均为0.0：

pred_mask=(rs[0].masks.masks[0].numpy()*255).astype("uint8")

#some code ommitted here

with open(pt_json,"r",encoding="utf-8") as f:
    dj=json.load(f)
grount_truth_mask=labelme.utils.shape_to_mask((dj['imageHeight'],dj['imageWidth']),dj['shapes'][0]['points'],shape_type=None,line_width=1,point_size=1)

pred_area=st.resize(pred_area,(640,480),order=0,preserve_range=True,anti_aliasing=False)
mask_area=st.resize(mask_area,(640,480),order=0,preserve_range=True,anti_aliasing=False)


pred_area=pred_area.ravel().copy()
mask_area=mask_area.ravel().copy()

pred_area=pred_area==1
mask_area=mask_area==1


intersection = np.sum(pred_area*mask_area)
union=pred_area.sum()+mask_area.sum()-intersection
iou=np.mean(intersection/union)

问题原因分析 #

• 变量名与赋值不匹配：代码中定义了pred_mask但后续使用pred_area，grount_truth_mask（拼写错误）对应后续的mask_area，存在未初始化变量或变量错位问题。
• 掩码二值化逻辑错误：YOLOv8的掩码乘以255后是uint8类型（值为0或255），但代码中用pred_area==1判断，导致预测区域全为False，IoU自然为0。
• 尺寸与坐标系未对齐：resize时目标尺寸顺序可能与掩码的宽高顺序不匹配，导致掩码错位无重叠。

修正后的代码 #

import numpy as np
import labelme.utils
import skimage.transform as st
import json

# 从YOLOv8推理结果提取预测掩码
pred_mask = (rs[0].masks.masks[0].numpy() * 255).astype("uint8")

# 加载LabelMe标注文件并生成真值掩码
with open(pt_json, "r", encoding="utf-8") as f:
    dj = json.load(f)
ground_truth_mask = labelme.utils.shape_to_mask(
    (dj['imageHeight'], dj['imageWidth']),
    dj['shapes'][0]['points'],
    shape_type=None,
    line_width=1,
    point_size=1
)

# 统一掩码尺寸：注意目标尺寸顺序为(高度, 宽度)，与shape_to_mask输入一致
target_size = (480, 640)
pred_resized = st.resize(pred_mask, target_size, order=0, preserve_range=True, anti_aliasing=False).astype("uint8")
gt_resized = st.resize(ground_truth_mask.astype("uint8"), target_size, order=0, preserve_range=True, anti_aliasing=False).astype("uint8")

# 转换为二值布尔数组：适配YOLOv8的0/255和LabelMe的0/1掩码格式
pred_binary = pred_resized > 0
gt_binary = gt_resized > 0

# 计算IoU，处理union为0的边界情况
intersection = np.logical_and(pred_binary, gt_binary).sum()
union = np.logical_or(pred_binary, gt_binary).sum()

iou = intersection / union if union != 0 else 0.0
print(f"IoU值: {iou}")

关键注意事项 #

• 变量一致性：确保掩码变量的定义与使用完全对应，修正拼写错误（如grount_truth_mask改为ground_truth_mask）。
• 二值化适配：用>0替代==1判断，兼容YOLOv8的255掩码值和LabelMe的1掩码值。
• 尺寸对齐：resize目标尺寸顺序保持为(高度, 宽度)，与LabelMe标注的imageHeight/imageWidth顺序一致，避免宽高颠倒。
• 边界处理：增加union为0的判断，避免除以0报错。

内容的提问来源于stack exchange，提问作者Anitha Anitha