Spaces:
Runtime error
Runtime error
| import math | |
| from dataclasses import dataclass | |
| SQRT2 = math.sqrt(2) | |
| class Circle: | |
| x: int | |
| y: int | |
| r: int | |
| class Corners: | |
| top_left: Circle | |
| top_right: Circle | |
| bottom_left: Circle | |
| bottom_right: Circle | |
| def from_opposite_corners(cls, top_left: Circle, bottom_right: Circle) -> "Corners": | |
| r_tl = top_left.r | |
| r_br = bottom_right.r | |
| r_max = max(r_tl, r_br) | |
| x_tl, y_tl = float(top_left.x), float(top_left.y) | |
| x_br, y_br = float(bottom_right.x), float(bottom_right.y) | |
| if r_tl < r_max: | |
| diff = r_max - r_tl | |
| dx = x_br - x_tl | |
| dy = y_br - y_tl | |
| dist = math.sqrt(dx * dx + dy * dy) | |
| if dist > 0: | |
| x_tl = x_tl + diff * dx / dist | |
| y_tl = y_tl + diff * dy / dist | |
| if r_br < r_max: | |
| diff = r_max - r_br | |
| dx = x_tl - x_br | |
| dy = y_tl - y_br | |
| dist = math.sqrt(dx * dx + dy * dy) | |
| if dist > 0: | |
| x_br = x_br + diff * dx / dist | |
| y_br = y_br + diff * dy / dist | |
| x_tl, y_tl = int(round(x_tl)), int(round(y_tl)) | |
| x_br, y_br = int(round(x_br)), int(round(y_br)) | |
| return cls( | |
| top_left=Circle(x_tl, y_tl, r_max), | |
| top_right=Circle(x_br, y_tl, r_max), | |
| bottom_left=Circle(x_tl, y_br, r_max), | |
| bottom_right=Circle(x_br, y_br, r_max), | |
| ) | |
| def from_boxes(cls, boxes: list[tuple[int, int, int, int]]) -> "Corners | None": | |
| if len(boxes) < 4: | |
| return None | |
| box_centers = [] | |
| for x1, y1, x2, y2 in boxes[:4]: | |
| cx = (x1 + x2) / 2 | |
| cy = (y1 + y2) / 2 | |
| box_centers.append((cx, cy, x1, y1, x2, y2)) | |
| sorted_by_y = sorted(box_centers, key=lambda b: b[1]) | |
| top_two = sorted(sorted_by_y[:2], key=lambda b: b[0]) | |
| bottom_two = sorted(sorted_by_y[2:], key=lambda b: b[0]) | |
| def box_to_circle(box_data: tuple, corner: str) -> Circle: | |
| _, _, x1, y1, x2, y2 = box_data | |
| r = (x2 - x1 + y2 - y1) // 2 | |
| if corner == "top_left": | |
| return Circle(x2, y2, r) | |
| elif corner == "top_right": | |
| return Circle(x1, y2, r) | |
| elif corner == "bottom_left": | |
| return Circle(x2, y1, r) | |
| else: | |
| return Circle(x1, y1, r) | |
| return cls( | |
| top_left=box_to_circle(top_two[0], "top_left"), | |
| top_right=box_to_circle(top_two[1], "top_right"), | |
| bottom_left=box_to_circle(bottom_two[0], "bottom_left"), | |
| bottom_right=box_to_circle(bottom_two[1], "bottom_right"), | |
| ) | |
| def to_dict(self) -> dict: | |
| return { | |
| "top_left": {"x": self.top_left.x, "y": self.top_left.y, "r": self.top_left.r}, | |
| "top_right": {"x": self.top_right.x, "y": self.top_right.y, "r": self.top_right.r}, | |
| "bottom_left": {"x": self.bottom_left.x, "y": self.bottom_left.y, "r": self.bottom_left.r}, | |
| "bottom_right": {"x": self.bottom_right.x, "y": self.bottom_right.y, "r": self.bottom_right.r}, | |
| } | |
| def get_crop_bounds(self) -> tuple[int, int, int, int]: | |
| x_left_tl = self.top_left.x - int(self.top_left.r / SQRT2) | |
| x_left_bl = self.bottom_left.x - int(self.bottom_left.r / SQRT2) | |
| x_left = max(x_left_tl, x_left_bl) | |
| x_right_tr = self.top_right.x + int(self.top_right.r / SQRT2) | |
| x_right_br = self.bottom_right.x + int(self.bottom_right.r / SQRT2) | |
| x_right = min(x_right_tr, x_right_br) | |
| y_top_tl = self.top_left.y - int(self.top_left.r / SQRT2) | |
| y_top_tr = self.top_right.y - int(self.top_right.r / SQRT2) | |
| y_top = max(y_top_tl, y_top_tr) | |
| y_bottom_bl = self.bottom_left.y + int(self.bottom_left.r / SQRT2) | |
| y_bottom_br = self.bottom_right.y + int(self.bottom_right.r / SQRT2) | |
| y_bottom = min(y_bottom_bl, y_bottom_br) | |
| return x_left, y_top, x_right, y_bottom | |