本文记录了在目标检测任务中根据目标的真实坐标和预测坐标计算 iou 交并比指标的代码。
文章目录
- 一、代码
一、代码
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimgdef 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_widthy_center = y_center * img_heightwidth = width * img_widthheight = height * img_heightx_min = x_center - width / 2y_min = y_center - height / 2x_max = x_center + width / 2y_max = y_center + height / 2return [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# 计算IoUiou = inter_area / union_areareturn ioudef 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_widthy_center = y_center * img_heightwidth = width * img_widthheight = height * img_heightx_min = x_center - width / 2y_min = y_center - height / 2x_max = x_center + width / 2y_max = y_center + height / 2return [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] # 示例预测标签# 计算IoUimg = 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)