【Python】2D/3D框IOU简单计算方法

算是破事水了哈哈哈

还是记录一下吧

万一能帮助到别人

一、2D框

思路+原理：

以下都以矩形框为例

• 首先，框必须有能确定4个顶点坐标的参数，我这里用的中心坐标+长宽。如果需要旋转，还需要旋转角度。下以逆时针旋转为例。

• 旋转的原理是先通过旋转角度构造旋转矩阵，然后左乘坐标进行旋转：

  • 注意这个情景下使用矩阵乘法时坐标得是这个格式：$[(x1,x2,x3,x4,...),(y1,y2,y3,y4,...)]$

  • 二维旋转都是绕原点旋转，如果希望绕某点旋转，得通过先平移到原点，再平移回去的方式完成（如果有其他方式请告诉我）。

  • 这个旋转函数不用自己写，我们直接使用shapely里的rotate完成。下面的代码只是给大家看看原理。

  def to_2d_rot_matrix(angle, degrees=False):
      """
      计算二维平面旋转矩阵的函数。
      angle: 旋转角度。
      degrees: 使用弧度制还是角度值。默认False，表示弧度制。
      """
      if degrees:
          angle *= np.pi / 180
      return [
          [np.cos(angle), -np.sin(angle)],
          [np.sin(angle), np.cos(angle)]
      ]
  
  def rot_2d(points, rot_mat, tx=0, ty=0):
      """
      进行二维旋转的函数。
      points: 点坐标，[(x1, y1), (x2, y2), ...]。
      rot_mat: 旋转矩阵。
      tx: x方向上的平移。
      ty: y方向上的平移。
      (tx, ty)也就是旋转中心。
      """
      points_x = np.array([p[0] for p in points]) - tx
      points_y = np.array([p[1] for p in points]) - ty
      
      if type(rot_mat) != np.array:
          rot_mat = np.array(rot_mat)   
          
      result = rot_mat @ np.array([points_x, points_y]) + np.array([[tx], [ty]])
      return result.T.reshape(-1,2)
  

• 剩下就是具体实现：

  • 使用shapely.geometry里的box函数（传左下右上2点坐标，共4个值）来创造Polygon对象。
  • 使用affinity里的rotate完成旋转（注意弧度制要把参数use_radians改成True）。
  • 使用超级无敌牛逼的Polygon.intersection方法计算交集，使用同样牛逼的union方法计算并集，然后相除得到IOU。

完整代码如下：

from shapely.geometry import box 
from shapely import affinity

def my_iou_2d(box1, box2):
    """
    两个box均为5元素list，5个元素分别是中心点xy坐标、箱子长宽和偏航角（弧度制）
    """
    result_xy = []
    for b in [box1, box2]:
        # 先解包获取两框中心坐标、长宽、偏航角
        x, y, l, w, yaw = b

        # 构造矩形
        poly = box(x - l/2, y - w/2, x + l/2, y + w/2)
        poly_rot = affinity.rotate(poly, yaw, use_radians=True)
        result_xy.append(poly_rot)

    # 计算xy平面面积重叠、z轴重叠
    poly1, poly2 = result_xy
    
    # 计算IOU
    return poly1.intersection(poly2).area / poly1.union(poly2).area

box_2d_1 = np.array([1,1,2,2,np.pi/2])
box_2d_2 = np.array([1,0.8,2,2,np.pi/3])
my_iou_2d(box_2d_1, box_2d_2)
# 0.6980890270283112

二、3D框

思路+原理：

以下都以长方体框为例

• 首先，框必须有能确定8个顶点坐标的参数，我这里用的中心坐标+长宽高。如果需要绕z轴旋转，还需要旋转角度（即偏航角），绕x、y轴旋转的情况太复杂了不讨论。下以逆时针旋转为例。
• 其实3D框的情景只是比2D框多了个z轴，所以只需要先顺着z轴反方向看两个3D框，把他们想成2D框，计算出xy平面上的交集，再乘z轴交集，就是总的交集啦！
• xy平面的交集和2D框计算方法一样的，而z轴交集也可以通过LineString.intersection方法实现。

代码如下：

from shapely.geometry import LineString, box 
from shapely import affinity

def my_iou_3d(box1, box2):
    """
    两个box均为7元素list，7个元素分别是中心点xyz坐标、箱子长宽高和偏航角（弧度制）
    """
    result_xy, result_z, result_v = [], [], []
    for b in [box1, box2]:
        # 先解包获取两框中心坐标、长宽高、偏航角
        x, y, z, l, w, h, yaw = b
        
        # 计算体积
        result_v.append(l * w * h)
        
        # 构造z轴
        ls = LineString([[0, z - h/2], [0, z + h/2]])
        result_z.append(ls)
        
        # 构造xy平面部分的矩形
        poly = box(x - l/2, y - w/2, x + l/2, y + w/2)
        poly_rot = affinity.rotate(poly, yaw, use_radians=True)
        result_xy.append(poly_rot)

    # 计算xy平面面积重叠、z轴重叠
    overlap_xy = result_xy[0].intersection(result_xy[1]).area
    overlap_z = result_z[0].intersection(result_z[1]).length
    
    # 计算IOU
    overlap_xyz = overlap_z * overlap_xy
    return overlap_xyz / (np.sum(result_v) - overlap_xyz)

box_3d_1 = np.array([1,1,1,2,2,2,np.pi/2])
box_3d_2 = np.array([1,0.8,0.8,2,2,2,np.pi/3])
my_iou_3d(box_3d_1, box_3d_2)
# 0.5872825723717148