Bounding Box を表現するクラスとIoU計算の関数

概要

• 物体検出のタスクでバウンディングボックスをコード内で表現するとき，単に数値4つのタプルを使ってると(x1,y1,x2,y2)や(x,y,w,h)などの形式のうちどれを使ってるかわからなくなって可読性下がる問題の解決策
• dataclasses.dataclass最強

コード

from dataclasses import dataclass


@dataclass
class BoundingBox:
    x1: float
    y1: float
    x2: float
    y2: float

    @property
    def cx(self):
        return self.x1 + self.width / 2.0

    @property
    def cy(self):
        return self.y1 + self.height / 2.0

    @property
    def width(self):
        return self.x2 - self.x1

    @property
    def height(self):
        return self.y2 - self.y1

    @property
    def area(self):
        return self.width * self.height


def compute_iou(bb1: BoundingBox, bb2: BoundingBox) -> float:
    x1 = max(bb1.x1, bb2.x1)
    y1 = max(bb1.y1, bb2.y1)
    x2 = min(bb1.x2, bb2.x2)
    y2 = min(bb1.y2, bb2.y2)

    inter_bb_w = max(0, x2 - x1)
    inter_bb_h = max(0, y2 - y1)
    inter_bb_area = inter_bb_w * inter_bb_h

    return inter_bb_area / (bb1.area + bb2.area - inter_bb_area)