wehub-resource-sync
123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
import copy
|
|
from typing import List, Optional
|
|
|
|
import numpy as np
|
|
from pydantic import BaseModel, field_validator, computed_field
|
|
import numbers
|
|
|
|
|
|
class PolygonBox(BaseModel):
|
|
polygon: List[List[float]]
|
|
confidence: Optional[float] = None
|
|
|
|
@field_validator("polygon", mode="before")
|
|
@classmethod
|
|
def convert_bbox_to_polygon(cls, value):
|
|
if isinstance(value, (list, tuple)) and len(value) == 4:
|
|
if all(isinstance(x, numbers.Number) for x in value):
|
|
value = [float(v) for v in value]
|
|
x_min, y_min, x_max, y_max = value
|
|
polygon = [
|
|
[x_min, y_min],
|
|
[x_max, y_min],
|
|
[x_max, y_max],
|
|
[x_min, y_max],
|
|
]
|
|
return polygon
|
|
elif all(
|
|
isinstance(point, (list, tuple)) and len(point) == 2 for point in value
|
|
):
|
|
value = [[float(v) for v in point] for point in value]
|
|
return value
|
|
elif isinstance(value, np.ndarray):
|
|
if value.shape == (4, 2):
|
|
return value.tolist()
|
|
|
|
raise ValueError(
|
|
f"Input must be either a bbox [x_min, y_min, x_max, y_max] or a polygon with 4 corners [(x,y), (x,y), (x,y), (x,y)]. All values must be numeric. You passed {value} of type {type(value)}. The first value is of type {type(value[0])}."
|
|
)
|
|
|
|
@property
|
|
def height(self):
|
|
return self.bbox[3] - self.bbox[1]
|
|
|
|
@property
|
|
def width(self):
|
|
return self.bbox[2] - self.bbox[0]
|
|
|
|
@property
|
|
def area(self):
|
|
return self.width * self.height
|
|
|
|
@computed_field
|
|
@property
|
|
def bbox(self) -> List[float]:
|
|
x_coords = [point[0] for point in self.polygon]
|
|
y_coords = [point[1] for point in self.polygon]
|
|
return [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
|
|
|
def rescale(self, processor_size, image_size):
|
|
# Point is in x, y format
|
|
page_width, page_height = processor_size
|
|
|
|
img_width, img_height = image_size
|
|
width_scaler = img_width / page_width
|
|
height_scaler = img_height / page_height
|
|
|
|
for corner in self.polygon:
|
|
corner[0] = int(corner[0] * width_scaler)
|
|
corner[1] = int(corner[1] * height_scaler)
|
|
|
|
def round(self, divisor):
|
|
for corner in self.polygon:
|
|
corner[0] = int(corner[0] / divisor) * divisor
|
|
corner[1] = int(corner[1] / divisor) * divisor
|
|
|
|
def fit_to_bounds(self, bounds):
|
|
new_corners = copy.deepcopy(self.polygon)
|
|
for corner in new_corners:
|
|
corner[0] = max(min(corner[0], bounds[2]), bounds[0])
|
|
corner[1] = max(min(corner[1], bounds[3]), bounds[1])
|
|
self.polygon = new_corners
|
|
|
|
def expand(self, x_margin: float, y_margin: float):
|
|
new_polygon = []
|
|
x_margin = x_margin * self.width
|
|
y_margin = y_margin * self.height
|
|
for idx, poly in enumerate(self.polygon):
|
|
if idx == 0:
|
|
new_polygon.append([int(poly[0] - x_margin), int(poly[1] - y_margin)])
|
|
elif idx == 1:
|
|
new_polygon.append([int(poly[0] + x_margin), int(poly[1] - y_margin)])
|
|
elif idx == 2:
|
|
new_polygon.append([int(poly[0] + x_margin), int(poly[1] + y_margin)])
|
|
elif idx == 3:
|
|
new_polygon.append([int(poly[0] - x_margin), int(poly[1] + y_margin)])
|
|
self.polygon = new_polygon
|
|
|
|
def intersection_area(self, other, x_margin=0, y_margin=0):
|
|
x_overlap = self.x_overlap(other, x_margin)
|
|
y_overlap = self.y_overlap(other, y_margin)
|
|
return x_overlap * y_overlap
|
|
|
|
def x_overlap(self, other, x_margin=0):
|
|
return max(
|
|
0,
|
|
min(self.bbox[2] + x_margin, other.bbox[2] + x_margin)
|
|
- max(self.bbox[0] - x_margin, other.bbox[0] - x_margin),
|
|
)
|
|
|
|
def y_overlap(self, other, y_margin=0):
|
|
return max(
|
|
0,
|
|
min(self.bbox[3] + y_margin, other.bbox[3] + y_margin)
|
|
- max(self.bbox[1] - y_margin, other.bbox[1] - y_margin),
|
|
)
|
|
|
|
@property
|
|
def center(self):
|
|
return [(self.bbox[0] + self.bbox[2]) / 2, (self.bbox[1] + self.bbox[3]) / 2]
|
|
|
|
def __hash__(self):
|
|
return hash(tuple(self.bbox))
|