概要
- 物体検出のタスクでバウンディングボックスをコード内で表現するとき,単に数値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)