Python计算目标检测中的IoU

前言

• 本文是个人使用PythonPython处理文件的电子笔记，由于水平有限，难免出现错漏，敬请批评改正。
• 更多精彩内容，可点击进入我的个人主页查看

前提条件

• 熟悉Python

相关介绍

• Python是一种跨平台的计算机程序设计语言。是一个高层次的结合了解释性、编译性、互动性和面向对象的脚本语言。最初被设计用于编写自动化脚本(shell)，随着版本的不断更新和语言新功能的添加，越多被用于独立的、大型项目的开发。

实验环境

• Python 3.x （面向对象的高级语言）

IoU概念

• loU，英文全称为Intersection of Union，中文称作交并比。对于目标检测识别而言它是一种用来衡量真实框与预测框的贴合程度的方法。
• 假设真实框为$A$，预测框为$B$，$A\cap B$表示真实框与预测框的交集，$A\cup B$表示真实框与预测框的并集loU具体的计算方式，如图所示。
  

代码实现

import cv2
import numpy as np
def CountIOU(RecA, RecB): # Rec = [xmin,ymin,xmax,ymax]
    # 获取交集区域的[xmin,ymin,xmax,ymax]
    x_A_and_B_min = max(RecA[0], RecB[0])
    y_A_and_B_min = max(RecA[1], RecB[1])
    x_A_and_B_max = min(RecA[2], RecB[2])
    y_A_and_B_max = min(RecA[3], RecB[3])
    # 计算交集部分面积, 当(xmax - xmin)为负时，说明A与B框无交集，直接置为0。 (ymax - ymin)同理。
    interArea = max(0, x_A_and_B_max - x_A_and_B_min) * max(0, y_A_and_B_max - y_A_and_B_min)
    # 计算A和B的面积
    RecA_Area = (RecA[2] - RecA[0]) * (RecA[3] - RecA[1]) # (xmax - xmin) * (ymax - ymin)
    RecB_Area = (RecB[2] - RecB[0]) * (RecB[3] - RecB[1]) # (xmax - xmin) * (ymax - ymin)
    # 计算IOU
    iou = interArea / (RecA_Area + RecB_Area - interArea) 
    return iou

def draw_rect(img):
    cv2.rectangle(img, (RecA[0],RecA[1]), (RecA[2],RecA[3]), (0, 0, 255), 2)
    cv2.rectangle(img, (RecB[0],RecB[1]), (RecB[2],RecB[3]), (255, 0, 0), 2)
    cv2.putText(img,'({},{})'.format(RecA[0],RecA[1]),(RecA[0]-20,RecA[1]-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,255,0),2)
    cv2.putText(img,'({},{})'.format(RecA[2],RecA[3]),(RecA[2]-50,RecA[3]+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,255,0),2)
    cv2.putText(img,'({},{})'.format(RecB[0],RecB[1]),(RecB[0]-20,RecB[1]-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,255,0),2)
    cv2.putText(img,'({},{})'.format(RecB[2],RecB[3]),(RecB[2]-50,RecB[3]+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,255,0),2)

if __name__=="__main__":
    img = cv2.imread("cat1.jpg") 
    img_res = img.copy()
    # A框[xmin,ymin,xmax,ymax]
    RecA = [80,90,330,336]
    # B框[xmin,ymin,xmax,ymax]
    RecB = [200,255,425,490]
    # 原始图片
    draw_rect(img)
    # 计算IOU
    IOU = CountIOU(RecA,RecB)
    # 结果图片
    draw_rect(img_res)
    cv2.putText(img_res,"IOU=%.2f"%IOU,(205, 300),cv2.FONT_HERSHEY_SIMPLEX,0.8,(255,255,255),2)

    cv2.imshow("img",img)
    cv2.imshow("img_res",img_res)
    cv2.waitKey()
    cv2.destroyAllWindows()

更多精彩内容，可点击进入Python日常小操作专栏或我的个人主页查看