目标检测任务中根据真实坐标和预测坐标计算IOU

本文记录了在目标检测任务中根据目标的真实坐标和预测坐标计算 iou 交并比指标的代码。

一、代码

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg


def calculate_iou(real_label, predicted_label, img_width, img_height):
    """
    计算交并比（IoU）。
    :param real_label: list, 真实标签 [class, x_center, y_center, width, height]
    :param predicted_label: list, 预测标签 [class, x_center, y_center, width, height]
    :param img_width: int, 图像的宽度
    :param img_height: int, 图像的高度
    :return: float, IoU值
    """

    # 解码标签 (x_center, y_center, width, height -> x_min, y_min, x_max, y_max)
    def decode_bbox(label, img_width, img_height):
        _, x_center, y_center, width, height = label
        # 将归一化坐标转换为像素值
        x_center = x_center * img_width
        y_center = y_center * img_height
        width = width * img_width
        height = height * img_height

        x_min = x_center - width / 2
        y_min = y_center - height / 2
        x_max = x_center + width / 2
        y_max = y_center + height / 2
        return [x_min, y_min, x_max, y_max]

    real_bbox = decode_bbox(real_label, img_width, img_height)
    predicted_bbox = decode_bbox(predicted_label, img_width, img_height)

    # 计算交集
    inter_x_min = max(real_bbox[0], predicted_bbox[0])
    inter_y_min = max(real_bbox[1], predicted_bbox[1])
    inter_x_max = min(real_bbox[2], predicted_bbox[2])
    inter_y_max = min(real_bbox[3], predicted_bbox[3])

    inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)

    # 计算并集
    real_area = (real_bbox[2] - real_bbox[0]) * (real_bbox[3] - real_bbox[1])
    predicted_area = (predicted_bbox[2] - predicted_bbox[0]) * (predicted_bbox[3] - predicted_bbox[1])

    union_area = real_area + predicted_area - inter_area

    # 避免除零错误
    if union_area == 0:
        return 0

    # 计算IoU
    iou = inter_area / union_area
    return iou


def plot_bboxes(image_path, real_label, predicted_label, real_color='green', pred_color='red', linewidth=2):
    """
    绘制真实标签和预测标签的边界框，并显示图片。
    :param image_path: str, 图片文件的路径
    :param real_label: list, 真实标签 [class, x_center, y_center, width, height]
    :param predicted_label: list, 预测标签 [class, x_center, y_center, width, height]
    :param real_color: str, 真实标签框的颜色
    :param pred_color: str, 预测标签框的颜色
    :param linewidth: int, 线宽
    """
    # 读取图像
    img = mpimg.imread(image_path)
    img_height, img_width = img.shape[:2]  # 获取图像的宽高

    # 解码真实和预测框
    def decode_bbox(label, img_width, img_height):
        _, x_center, y_center, width, height = label
        # 将归一化坐标转换为像素值
        x_center = x_center * img_width
        y_center = y_center * img_height
        width = width * img_width
        height = height * img_height

        x_min = x_center - width / 2
        y_min = y_center - height / 2
        x_max = x_center + width / 2
        y_max = y_center + height / 2
        return [x_min, y_min, x_max, y_max]

    real_bbox = decode_bbox(real_label, img_width, img_height)
    predicted_bbox = decode_bbox(predicted_label, img_width, img_height)

    # 创建图像和坐标轴
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.imshow(img)

    # 绘制真实标签框
    real_rect = plt.Rectangle((real_bbox[0], real_bbox[1]),
                              real_bbox[2] - real_bbox[0],
                              real_bbox[3] - real_bbox[1],
                              edgecolor=real_color, facecolor='none', linewidth=linewidth)
    ax.add_patch(real_rect)

    # 绘制预测标签框
    predicted_rect = plt.Rectangle((predicted_bbox[0], predicted_bbox[1]),
                                   predicted_bbox[2] - predicted_bbox[0],
                                   predicted_bbox[3] - predicted_bbox[1],
                                   edgecolor=pred_color, facecolor='none', linewidth=linewidth)
    ax.add_patch(predicted_rect)

    # 设置图像边界
    ax.set_xlim(0, img_width)
    ax.set_ylim(0, img_height)
    ax.invert_yaxis()  # 坐标系与图像坐标一致（左上角为原点）

    # 隐藏坐标轴和数字
    ax.axis('off')

    # 不显示标题
    plt.title('')

    # 不显示图例
    # plt.legend()  # 移除这行以不显示图例

    # 调整图像大小以填充整个画布
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

    # 显示图像
    plt.show()




# 示例用法
if __name__ == "__main__":
    image_path = 'D:\\images_origin\\0003501.jpg'  # 替换为实际图像路径
    real_label = [0, 0.653302, 0.643799, 0.693396, 0.712402]  # [class, x_center, y_center, width, height]
    predicted_label = [0, 0.658956, 0.658806, 0.682088, 0.673066]  # 示例预测标签

    # 计算IoU
    img = mpimg.imread(image_path)
    img_height, img_width = img.shape[:2]
    iou = calculate_iou(real_label, predicted_label, img_width, img_height)
    print(f"IoU: {iou:.4f}")

    # 绘制边界框
    plot_bboxes(image_path, real_label, predicted_label, real_color='green', pred_color='yellow', linewidth=2)