56 lines
1.4 KiB
Python
56 lines
1.4 KiB
Python
import dataclasses
|
|
import numpy as np
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class BoundingBox:
|
|
x_min: float
|
|
y_min: float
|
|
x_max: float
|
|
y_max: float
|
|
label: str
|
|
mask_uid: str = ""
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"x_min": self.x_min,
|
|
"y_min": self.y_min,
|
|
"x_max": self.x_max,
|
|
"y_max": self.y_max,
|
|
"label": self.label,
|
|
"mask_uid": self.mask_uid
|
|
}
|
|
|
|
@property
|
|
def center(self):
|
|
return np.array([(self.x_min + self.x_max) / 2, (self.y_min + self.y_max) / 2])
|
|
|
|
def distance_to(self, point: np.ndarray):
|
|
return np.linalg.norm(self.center - point)
|
|
|
|
def contains(self, point: np.ndarray):
|
|
return self.x_min <= point[0] <= self.x_max and self.y_min <= point[1] <= self.y_max
|
|
|
|
|
|
def get_mask_bounding_box(mask, label: str):
|
|
where = np.where(mask)
|
|
x_min = np.min(where[1])
|
|
y_min = np.min(where[0])
|
|
x_max = np.max(where[1])
|
|
y_max = np.max(where[0])
|
|
return BoundingBox(
|
|
x_min / mask.shape[1],
|
|
y_min / mask.shape[0],
|
|
x_max / mask.shape[1],
|
|
y_max / mask.shape[0],
|
|
label
|
|
)
|
|
|
|
|
|
def get_bounding_boxes(masks, labels):
|
|
bounding_boxes = []
|
|
for mask, label in zip(masks, labels):
|
|
bounding_box = get_mask_bounding_box(mask, label)
|
|
bounding_boxes.append(bounding_box)
|
|
return bounding_boxes
|